/* Copyright (C) 2015 Free Software Foundation, Inc.
   This file is part of the GNU C Library.

   The GNU C Library is free software; you can redistribute it and/or
   modify it under the terms of the GNU Lesser General Public
   License as published by the Free Software Foundation; either
   version 2.1 of the License, or (at your option) any later version.

   The GNU C Library is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
   Lesser General Public License for more details.

   You should have received a copy of the GNU Lesser General Public
   License along with the GNU C Library; if not, see
   <http://www.gnu.org/licenses/>.  */

#include <errno.h>
#include <limits.h>
#include <sched.h>
#include <stdbool.h>
#include <stdio.h>

static int
find_max_cpu (void)
{
  /* We need to start at 64 because otherwise, CPU_ALLOC
     over-allocates, and and we do not see all bits returned by the
     kernel.  */
  for (int num_cpus = 64; num_cpus <= INT_MAX / 2; num_cpus *= 2)
    {
      cpu_set_t *set = CPU_ALLOC (num_cpus);
      size_t size = CPU_ALLOC_SIZE (num_cpus);
      printf ("allocating set for %d CPUs, size %zu\n", num_cpus, size);

      if (set == NULL)
	{
	  printf ("CPU_ALLOC(%d) failed\n", num_cpus);
	  return -1;
	}
      if (sched_getaffinity (0, size, set) == 0)
	{
	  for (int i = num_cpus - 1; i >= 0; --i)
	    {
	      if (CPU_ISSET_S (i, size, set))
		{
		  CPU_FREE (set);
		  return i;
		}
	    }
	  printf ("Could not find active CPU in set of size %d\n", num_cpus);
	  CPU_FREE (set);
	  return -1;
	}
      if (errno != EINVAL)
	{
	  printf ("sched_getaffinity for %d CPUs: %m\n", num_cpus);
	  return -1;
	}
      CPU_FREE (set);
    }
  puts ("Cannot find maximum CPU number");
  return -1;
}

static int
find_flippable_cpu (const cpu_set_t *set, size_t size)
{
  if (CPU_COUNT_S (size, set) < 2)
    return -1;

  int i;
  for (i = 0; !CPU_ISSET_S (i, size, set); ++i)
    ;
  return i;
}

static bool
test_size (size_t size)
{
  cpu_set_t *set = CPU_ALLOC (size);
  cpu_set_t *set2 = CPU_ALLOC (size);
  cpu_set_t *active_cpu_set = CPU_ALLOC (size);
  int active_cpu;
  
  if (set == NULL || set2 == NULL || active_cpu_set == NULL)
    {
      printf ("size %zu: CPU_ALLOC failed\n", size);
      return false;
    }
  size = CPU_ALLOC_SIZE (size);
  
  if (sched_getaffinity (0, size, set) < 0)
    {
      printf ("size %zu: sched_getaffinity: %m\n", size);
      return false;
    }
  active_cpu = find_flippable_cpu (set, size);
  if (active_cpu < 0)
    {
      printf ("size %zu: not enough CPUs to test affinity\n", size);
      return true;
    }
  if (sched_setaffinity (0, size, set) < 0)
    {
      printf ("size %zu: sched_setaffinity: %m\n", size);
      return true;
    }

  CPU_ZERO_S (size, active_cpu_set);
  CPU_SET_S (active_cpu, size, active_cpu_set);
  if (sched_setaffinity (0, size, active_cpu_set) < 0)
    {
      printf ("size %zu: sched_setaffinity (2): %m\n", size);
      return false;
    }
  if (sched_getaffinity (0, size, set2) < 0)
    {
      printf ("size %zu: sched_getaffinity (2): %m\n", size);
      return false;
    }
  if (!CPU_EQUAL_S(size, active_cpu_set, set2))
    {
      printf ("size %zu: CPU sets do not match\n", size);
      return false;
    }
  if (sched_setaffinity (0, size, set) < 0)
    {
      printf ("size %zu: sched_setaffinity (3): %m\n", size);
      return false;
    }
  if (sched_getaffinity (0, size, set2) < 0)
    {
      printf ("size %zu: sched_getaffinity (3): %m\n", size);
      return false;
    }
  if (!CPU_EQUAL_S(size, set, set2))
    {
      printf ("size %zu: CPU sets do not match (2)\n", size);
      return false;
    }

  CPU_FREE (set);
  CPU_FREE (set2);
  CPU_FREE (active_cpu_set);

  return true;
}

static bool
shrink_set (int max_cpu, bool *shrunk)
{
  cpu_set_t *set = CPU_ALLOC (max_cpu + 1);
  size_t size = CPU_ALLOC_SIZE (max_cpu + 1);
    
  if (set == NULL)
    {
      printf ("CPU_ALLOC(%d) failed\n", max_cpu);
      return false;
    }

  if (sched_getaffinity (0, size, set) < 0)
    {
      printf ("sched_getaffinity: %m\n");
      CPU_FREE (set);
      return false;
    }

  *shrunk = CPU_ISSET_S (0, size, set) && CPU_ISSET_S (0, size, set);
  if (!*shrunk)
    {
      printf ("Cannot shrink CPU set size because CPUs 0 and 1 are missing\n");
      CPU_FREE (set);
      return true;
    }
  CPU_ZERO_S (size, set);
  CPU_SET_S (0, size, set);
  CPU_SET_S (1, size, set);
  if (sched_setaffinity (0, size, set) < 0)
    {
      printf ("sched_setaffinity: %m\n");
      CPU_FREE (set);
      return false;
    }
  max_cpu = find_max_cpu ();
  if (max_cpu != 1)
    {
      printf ("Set shrinking unsuccessful, maximum CPU still %d\n", max_cpu);
      CPU_FREE (set);
      return false;
    }
  
  CPU_FREE (set);
  return true;
}

static int
do_test (void)
{
  printf ("__NCPUBITS: %d\nsizeof (__cpu_mask): %zu\n",
	  (int)__NCPUBITS, sizeof (__cpu_mask));
  
  cpu_set_t set;
  if (sched_getaffinity (0, sizeof (set), &set) < 0)
    {
      if (errno == ENOSYS)
	{
	  puts ("sched_getaffinity not supported");
	  return 0;
	}
      printf ("sched_getaffinity: %m\n");
      return 1;
    }

  int max_cpu = find_max_cpu ();
  printf ("Detected maximum CPU number: %d\n", max_cpu);
  if (max_cpu < 0)
    return 1;

  if (max_cpu <= 1024 && !test_size (1024))
    return 1;
  if (max_cpu <= 2 && !test_size (2))
    return 1;
  if (max_cpu <= 1024 * 1024 && !test_size (1024 * 1024))
    return 1;

  bool shrunk;
  if (!shrink_set (max_cpu, &shrunk))
      return 1;
  if (shrunk)
    {
      if (!(test_size (1024) && test_size (2) && test_size (1024 * 1024)))
	return 1;
    }

  return 0;
}

#define TEST_FUNCTION do_test ()
#include "../test-skeleton.c"