string: add __mtag_tag_region

Add optimized __mtag_tag_region(dst, len) operation to AOR. It tags the
given memory region according to the tag of the dst pointer and returns
dst. It requires MTE support. The memory remains untagged if tagging is
not enabled for it. The dst must be 16 bytes aligned and len must be a
multiple of 16.

4 files changed, 251 insertions, 0 deletions

diff --git a/string/Dir.mk b/string/Dir.mk
index 7817357..863b27e 100644
--- a/string/Dir.mk
+++ b/string/Dir.mk
@@ -29,6 +29,7 @@ string-tests := \
 	build/bin/test/memchr \
 	build/bin/test/memrchr \
 	build/bin/test/memcmp \
+	build/bin/test/__mtag_tag_region \
 	build/bin/test/strcpy \
 	build/bin/test/stpcpy \
 	build/bin/test/strcmp \
diff --git a/string/aarch64/__mtag_tag_region.S b/string/aarch64/__mtag_tag_region.S
new file mode 100644
index 0000000..84339f7
--- /dev/null
+++ b/string/aarch64/__mtag_tag_region.S
@@ -0,0 +1,100 @@
+/*
+ * __mtag_tag_region - tag memory
+ *
+ * Copyright (c) 2021, Arm Limited.
+ * SPDX-License-Identifier: MIT
+ */
+
+/* Assumptions:
+ *
+ * ARMv8-a, AArch64, MTE, LP64 ABI.
+ *
+ * Interface contract:
+ * Address is 16 byte aligned and size is multiple of 16.
+ * Returns the passed pointer.
+ * The memory region may remain untagged if tagging is not enabled.
+ */
+
+#include "../asmdefs.h"
+
+#if __ARM_FEATURE_MEMORY_TAGGING
+
+#define dstin	x0
+#define count	x1
+#define dst	x2
+#define dstend	x3
+#define tmp	x4
+#define zva_val	x4
+
+ENTRY (__mtag_tag_region)
+	PTR_ARG (0)
+	SIZE_ARG (1)
+
+	add	dstend, dstin, count
+
+	cmp	count, 96
+	b.hi	L(set_long)
+
+	tbnz	count, 6, L(set96)
+
+	/* Set 0, 16, 32, or 48 bytes.  */
+	lsr	tmp, count, 5
+	add	tmp, dstin, tmp, lsl 4
+	cbz     count, L(end)
+	stg	dstin, [dstin]
+	stg	dstin, [tmp]
+	stg	dstin, [dstend, -16]
+L(end):
+	ret
+
+	.p2align 4
+	/* Set 64..96 bytes.  Write 64 bytes from the start and
+	   32 bytes from the end.  */
+L(set96):
+	st2g	dstin, [dstin]
+	st2g	dstin, [dstin, 32]
+	st2g	dstin, [dstend, -32]
+	ret
+
+	.p2align 4
+	/* Size is > 96 bytes.  */
+L(set_long):
+	cmp	count, 160
+	b.lo	L(no_zva)
+
+#ifndef SKIP_ZVA_CHECK
+	mrs	zva_val, dczid_el0
+	and	zva_val, zva_val, 31
+	cmp	zva_val, 4		/* ZVA size is 64 bytes.  */
+	b.ne	L(no_zva)
+#endif
+	st2g	dstin, [dstin]
+	st2g	dstin, [dstin, 32]
+	bic	dst, dstin, 63
+	sub	count, dstend, dst	/* Count is now 64 too large.  */
+	sub	count, count, 128	/* Adjust count and bias for loop.  */
+
+	.p2align 4
+L(zva_loop):
+	add	dst, dst, 64
+	dc	gva, dst
+	subs	count, count, 64
+	b.hi	L(zva_loop)
+	st2g	dstin, [dstend, -64]
+	st2g	dstin, [dstend, -32]
+	ret
+
+L(no_zva):
+	sub	dst, dstin, 32		/* Dst is biased by -32.  */
+	sub	count, count, 64	/* Adjust count for loop.  */
+L(no_zva_loop):
+	st2g	dstin, [dst, 32]
+	st2g	dstin, [dst, 64]!
+	subs	count, count, 64
+	b.hi	L(no_zva_loop)
+	st2g	dstin, [dstend, -64]
+	st2g	dstin, [dstend, -32]
+	ret
+
+END (__mtag_tag_region)
+#endif
diff --git a/string/include/stringlib.h b/string/include/stringlib.h
index 841a7bb..b3aa118 100644
--- a/string/include/stringlib.h
+++ b/string/include/stringlib.h
@@ -54,6 +54,9 @@ size_t __strlen_aarch64_sve (const char *);
 size_t __strnlen_aarch64_sve (const char *, size_t);
 int __strncmp_aarch64_sve (const char *, const char *, size_t);
 # endif
+# if __ARM_FEATURE_MEMORY_TAGGING
+void *__mtag_tag_region (void *, size_t);
+# endif
 #elif __arm__
 void *__memcpy_arm (void *__restrict, const void *__restrict, size_t);
 void *__memset_arm (void *, int, size_t);
diff --git a/string/test/__mtag_tag_region.c b/string/test/__mtag_tag_region.c
new file mode 100644
index 0000000..d8c02d9
--- /dev/null
+++ b/string/test/__mtag_tag_region.c
@@ -0,0 +1,147 @@
+/*
+ * __mtag_tag_region test.
+ *
+ * Copyright (c) 2021, Arm Limited.
+ * SPDX-License-Identifier: MIT
+ */
+
+#if __ARM_FEATURE_MEMORY_TAGGING && WANT_MTE_TEST
+#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include "mte.h"
+#include "stringlib.h"
+#include "stringtest.h"
+
+static void
+mtag_quoteat (const char *prefix, void *p, int len, int at)
+{
+  /* Print tag, untag and quote the context.  */
+  printf ("location: %p\n", __arm_mte_get_tag ((char *) p + at));
+  untag_buffer (p, len, 1);
+  p = untag_pointer (p);
+  quoteat (prefix, p, len, at);
+}
+
+#define F(x) {#x, x},
+
+static const struct fun
+{
+  const char *name;
+  void *(*fun) (void *s, size_t n);
+} funtab[] = {
+// clang-format off
+#if __aarch64__
+  F(__mtag_tag_region)
+#endif
+  {0, 0}
+  // clang-format on
+};
+#undef F
+
+#define A 64
+#define LEN 250000
+static unsigned char *sbuf;
+
+static void *
+alignup (void *p)
+{
+  return (void *) (((uintptr_t) p + A - 1) & -A);
+}
+
+static void
+test (const struct fun *fun, int salign, int len)
+{
+  unsigned char *src = alignup (sbuf);
+  unsigned char *s = src + salign;
+  void *p;
+  int i;
+
+  if (err_count >= ERR_LIMIT)
+    return;
+  if (len > LEN || salign >= A)
+    abort ();
+  for (i = 0; i < len + 2 * A; i++)
+    src[i] = '?';
+  for (i = 0; i < len; i++)
+    s[i] = 'a';
+
+  src = tag_buffer (src, len + 2 * A, 1);
+  s = src + salign;
+  /* Use different tag.  */
+  s = __arm_mte_increment_tag (s, 1);
+  p = fun->fun (s, len);
+
+  if (p != s)
+    ERR ("%s(%p,..) returned %p\n", fun->name, s, p);
+
+  for (i = 0; i < salign; i++)
+    {
+      if (src[i] != '?')
+	{
+	  ERR ("%s(align %d, %d) failed\n", fun->name, salign, len);
+	  mtag_quoteat ("got head", src, len + 2 * A, i);
+	  return;
+	}
+    }
+
+  for (; i < salign + len; i++)
+    {
+      if (s[i - salign] != 'a')
+	{
+	  ERR ("%s(align %d, %d) failed\n", fun->name, salign, len);
+	  mtag_quoteat ("got body", src, len + 2 * A, i);
+	  return;
+	}
+    }
+
+  for (; i < len + 2 * A; i++)
+    {
+      if (src[i] != '?')
+	{
+	  ERR ("%s(align %d, %d) failed\n", fun->name, salign, len);
+	  mtag_quoteat ("got tail", src, len + 2 * A, i);
+	  return;
+	}
+    }
+
+  untag_buffer (src, len + 2 * A, 1);
+}
+
+int
+main ()
+{
+  if (!mte_enabled ())
+    return 0;
+
+  sbuf = mte_mmap (LEN + 3 * A);
+  int r = 0;
+  for (int i = 0; funtab[i].name; i++)
+    {
+      err_count = 0;
+      for (int s = 0; s < A; s += 16)
+	{
+	  int n;
+	  for (n = 0; n < 200; n += 16)
+	    {
+	      test (funtab + i, s, n);
+	    }
+	  for (; n < LEN; n *= 2)
+	    {
+	      test (funtab + i, s, n);
+	    }
+	}
+      printf ("%s %s\n", err_count ? "FAIL" : "PASS", funtab[i].name);
+      if (err_count)
+	r = -1;
+    }
+  return r;
+}
+#else
+int
+main ()
+{
+  return 0;
+}
+#endif