#include <stdio.h>

#define NODES	3
#define ZONES	2

enum zone_stat_item {
	NUMA_HIT,
	NUMA_MISS,
	NUMA_LOCAL,
	NUMA_OTHER,
	NUMA_FOREIGN,
	ZONE_STAT_ITEMS
};

struct node;
struct zone {
	struct node *zone_pgdat;
	int node;
	int vm_stat[ZONE_STAT_ITEMS];
};

struct node {
	int nid;
	struct zone node_zones[ZONES];
};

typedef int gfp_t;
#define __GFP_OTHER_NODE	1
#define unlikely(x) (x)

int local_nid;
int numa_node_id()
{
	return local_nid;
}

static void __inc_zone_state(struct zone *z, enum zone_stat_item stat)
{
	z->vm_stat[stat]++;
}

// before b9f00e147f27
void zone_statistics1(struct zone *preferred_zone, struct zone *z, gfp_t flags)
{
	if (z->zone_pgdat == preferred_zone->zone_pgdat) {
		__inc_zone_state(z, NUMA_HIT);
	} else {
		__inc_zone_state(z, NUMA_MISS);
		__inc_zone_state(preferred_zone, NUMA_FOREIGN);
	}

	if (z->node == ((flags & __GFP_OTHER_NODE) ?
		preferred_zone->node : numa_node_id()))
                __inc_zone_state(z, NUMA_LOCAL);
        else
                __inc_zone_state(z, NUMA_OTHER);
}

// after b9f00e147f27
void zone_statistics2(struct zone *preferred_zone, struct zone *z, gfp_t flags)
{
        int local_nid = numa_node_id();
        enum zone_stat_item local_stat = NUMA_LOCAL;

        if (unlikely(flags & __GFP_OTHER_NODE)) {
                local_stat = NUMA_OTHER;
                local_nid = preferred_zone->node;
        }

        if (z->node == local_nid) {
                __inc_zone_state(z, NUMA_HIT);
                __inc_zone_state(z, local_stat);
        } else {
                __inc_zone_state(z, NUMA_MISS);
                __inc_zone_state(preferred_zone, NUMA_FOREIGN);
        }
}

// after Jia He's patch
static inline void zone_statistics3(struct zone *preferred_zone, struct zone *z,
                                                                gfp_t flags)
{
        int local_nid = numa_node_id();
        enum zone_stat_item local_stat = NUMA_LOCAL;

        if (unlikely(flags & __GFP_OTHER_NODE)) {
                local_stat = NUMA_OTHER;
                local_nid = preferred_zone->node;
        }

        if (z->node == local_nid) {
                __inc_zone_state(z, NUMA_HIT);
                __inc_zone_state(z, local_stat);
        } else if (z->node == preferred_zone->node) {
                __inc_zone_state(z, NUMA_HIT);
                __inc_zone_state(z, NUMA_OTHER);
        } else {
                __inc_zone_state(z, NUMA_MISS);
                __inc_zone_state(preferred_zone, NUMA_FOREIGN);
        }
}


void init_node(struct node * n, int nid)
{
	n->nid = nid;
	for (int i = 0; i < ZONES; i++) {
		struct zone * z = &n->node_zones[i];

		z->zone_pgdat = n;
		z->node = nid;
		for (int j = 0; j < ZONE_STAT_ITEMS; j++)
			z->vm_stat[j] = 0;
	}
}

struct node nodes1[NODES];
struct node nodes2[NODES];

void check_stats()
{
	for (int i = 0; i < NODES; i++)
		for (int j = 0; j < ZONES; j++)
			for (int k = 0; k < ZONE_STAT_ITEMS; k++)
				assert (nodes1[i].node_zones[j].vm_stat[k]
				     == nodes2[i].node_zones[j].vm_stat[k]);

}

int main(int argc, char *argv[])
{

	for (int i = 0; i < NODES; i++) {
		init_node(&nodes1[i], i);
		init_node(&nodes2[i], i);
	}

	local_nid = ((unsigned int)argv[1]) % NODES;
	int zone_nid = ((unsigned int)argv[2]) % NODES;
	int pzone_nid = ((unsigned int)argv[3]) % NODES;
	int zid = ((unsigned int)argv[4]) % ZONES;
	int pzid = ((unsigned int)argv[5]) % ZONES;

	/* we should not allocate from higher than preferred zone */
	if (zid > pzid)
		zid = pzid;

	gfp_t flags = ((unsigned int)argv[6]) & __GFP_OTHER_NODE;

	zone_statistics1(&nodes1[pzone_nid].node_zones[pzid],
			 &nodes1[zone_nid].node_zones[zid],
			 flags);

	zone_statistics3(&nodes2[pzone_nid].node_zones[pzid],
			 &nodes2[zone_nid].node_zones[zid],
			 flags);

	check_stats();
}