From d9bca644260847f3b325f4d0e12f0360ed1c0dac Mon Sep 17 00:00:00 2001
From: Javier Blazquez <jblazquez@riotgames.com>
Date: Tue, 9 Apr 2024 22:31:10 -0700
Subject: [PATCH] Add Windows ARM64 support (#656)

* Add Windows ARM64 support

* Add Windows ARM64/ARM64EC CI workflows
---
 .github/workflows/main.yml  | 23 +++++++++++++++++++++++
 src/snmalloc/aal/aal.h      |  9 ++++++---
 src/snmalloc/aal/aal_arm.h  |  2 +-
 src/snmalloc/ds_core/bits.h | 13 +++++++++++--
 4 files changed, 41 insertions(+), 6 deletions(-)

diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 0661ae414..bcb6b974f 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -388,6 +388,28 @@ jobs:
             build-type: Debug
             arch: x64
             toolchain: ""
+          - os: windows-2022
+            build-type: Release
+            arch: ARM64
+            toolchain: ""
+            build-only: yes
+          - os: windows-2022
+            build-type: Debug
+            arch: ARM64
+            toolchain: ""
+            build-only: yes
+          - os: windows-2022
+            build-type: Release
+            arch: ARM64EC
+            toolchain: ""
+            extra-cmake-flags: -DCMAKE_SYSTEM_VERSION="10.0.22621.0"
+            build-only: yes
+          - os: windows-2022
+            build-type: Debug
+            arch: ARM64EC
+            toolchain: ""
+            extra-cmake-flags: -DCMAKE_SYSTEM_VERSION="10.0.22621.0"
+            build-only: yes
         
       # Don't abort runners if a single one fails
       fail-fast: false
@@ -403,6 +425,7 @@ jobs:
       run: cmake --build ${{github.workspace}}/build -- /m /p:Configuration=${{ matrix.build-type }}
       # Run the tests.
     - name: Test
+      if: ${{ matrix.build-only != 'yes' }}
       working-directory: ${{ github.workspace }}/build
       run: ctest -j 2 --interactive-debug-mode 0 --output-on-failure -C ${{ matrix.build-type }} --timeout 400
       timeout-minutes: 20
diff --git a/src/snmalloc/aal/aal.h b/src/snmalloc/aal/aal.h
index 5014a7296..e1763995e 100644
--- a/src/snmalloc/aal/aal.h
+++ b/src/snmalloc/aal/aal.h
@@ -14,9 +14,11 @@
 #include <cstdint>
 #include <utility>
 
-#if defined(__i386__) || defined(_M_IX86) || defined(_X86_) || \
+#if ( \
+  defined(__i386__) || defined(_M_IX86) || defined(_X86_) || \
   defined(__amd64__) || defined(__x86_64__) || defined(_M_X64) || \
-  defined(_M_AMD64)
+  defined(_M_AMD64)) && \
+  !defined(_M_ARM64EC)
 #  if defined(SNMALLOC_SGX)
 #    define PLATFORM_IS_X86_SGX
 #    define SNMALLOC_NO_AAL_BUILTINS
@@ -25,7 +27,8 @@
 #  endif
 #endif
 
-#if defined(__arm__) || defined(__aarch64__)
+#if defined(__arm__) || defined(__aarch64__) || defined(_M_ARM64) || \
+  defined(_M_ARM64EC)
 #  define PLATFORM_IS_ARM
 #endif
 
diff --git a/src/snmalloc/aal/aal_arm.h b/src/snmalloc/aal/aal_arm.h
index 8b5a07e53..fef3aed37 100644
--- a/src/snmalloc/aal/aal_arm.h
+++ b/src/snmalloc/aal/aal_arm.h
@@ -1,6 +1,6 @@
 #pragma once
 
-#if defined(__aarch64__)
+#if defined(__aarch64__) || defined(_M_ARM64) || defined(_M_ARM64EC)
 #  define SNMALLOC_VA_BITS_64
 #  ifdef _MSC_VER
 #    include <arm64_neon.h>
diff --git a/src/snmalloc/ds_core/bits.h b/src/snmalloc/ds_core/bits.h
index b82ee846e..ac1c7e19b 100644
--- a/src/snmalloc/ds_core/bits.h
+++ b/src/snmalloc/ds_core/bits.h
@@ -158,7 +158,11 @@ namespace snmalloc
       SNMALLOC_ASSERT(x != 0); // Calling with 0 is UB on some implementations
 
 #if defined(_MSC_VER) && !defined(__clang__)
-#  ifdef _WIN64
+#  if defined(_M_ARM64) || defined(_M_ARM64EC)
+      unsigned long n = 0;
+      _BitScanForward64(&n, static_cast<unsigned __int64>(x));
+      return static_cast<size_t>(n);
+#  elif defined(_WIN64)
       return _tzcnt_u64(static_cast<unsigned __int64>(x));
 #  else
       return _tzcnt_u32(static_cast<unsigned int>(x));
@@ -203,7 +207,12 @@ namespace snmalloc
       overflow = __builtin_mul_overflow(x, y, &prod);
       return prod;
 #elif defined(_MSC_VER)
-#  ifdef _WIN64
+#  if defined(_M_ARM64) || defined(_M_ARM64EC)
+      size_t high_prod = __umulh(x, y);
+      size_t prod = x * y;
+      overflow = high_prod != 0;
+      return prod;
+#  elif defined(_WIN64)
       size_t high_prod;
       size_t prod = _umul128(x, y, &high_prod);
       overflow = high_prod != 0;