#include <stddef.h>
#include <stdint.h>
#include <sys/mman.h>
#include <sys/prctl.h>
#include <arm_acle.h>

int __memcmp_aarch64 (const void *, const void *, size_t);

#define PR_SET_TAGGED_ADDR_CTRL 55
#define PR_TAGGED_ADDR_ENABLE (1UL << 0)
#define PR_MTE_TCF_SHIFT 1
#define PR_MTE_TCF_SYNC (1UL << PR_MTE_TCF_SHIFT)
#define PR_MTE_TAG_SHIFT 3
#define PROT_MTE 0x20
#define MTE_GRANULE_SIZE 16

void *
alignup_mte (void *p)
{
  return (void *) (((uintptr_t) p + MTE_GRANULE_SIZE - 1)
		   & ~(MTE_GRANULE_SIZE - 1));
}

void *
aligndown_mte (void *p)
{
  return (void *) ((uintptr_t) p & ~(MTE_GRANULE_SIZE - 1));
}

void
tag_buffer_helper (void *p, int len)
{
  char *ptr = p;
  char *end = alignup_mte (ptr + len);
  ptr = aligndown_mte (p);
  for (; ptr < end; ptr += MTE_GRANULE_SIZE)
    {
      __arm_mte_set_tag (ptr);
    }
}

void *
tag_buffer (void *p, int len)
{
  p = __arm_mte_increment_tag (p, 1);
  tag_buffer_helper (p, len);
  return p;
}

int main (void)
{
  int r = prctl (PR_SET_TAGGED_ADDR_CTRL, PR_TAGGED_ADDR_ENABLE | PR_MTE_TCF_SYNC | (0xfffe << PR_MTE_TAG_SHIFT), 0, 0, 0);
  if (r < 0) return -1;
  char *src1 = mmap (NULL, 4096, PROT_READ | PROT_WRITE | PROT_MTE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
  char *src2 = mmap (NULL, 4096, PROT_READ | PROT_WRITE | PROT_MTE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
  if (src1 == MAP_FAILED) return -1;
  if (src2 == MAP_FAILED) return -1;
  char *s1 = src1;
  char *s2 = src2 + 15;
  for (int i = 0; i < 250; i++)
    src1[i] = src2[i] = '?';
  for (int i = 0; i < 200; i++)
    s1[i] = s2[i] = 'a' + i % 23;
  s1 = tag_buffer (s1, 200);
  s2 = tag_buffer (s2, 200);
  __memcmp_aarch64(s1, s2, 200);
  return 0;
}