From a393b5412fad1935b97c4d1067b83f21e73b6de1 Mon Sep 17 00:00:00 2001
From: Jacek Sieka <jacek@status.im>
Date: Thu, 26 Sep 2024 17:47:54 +0200
Subject: [PATCH] assign2: compile-time assignment support

---
 stew/assign2.nim       | 57 ++++++++++++++++++++++++++----------------
 tests/test_assign2.nim |  8 ++++++
 2 files changed, 44 insertions(+), 21 deletions(-)

diff --git a/stew/assign2.nim b/stew/assign2.nim
index 5b28326e..b85f6d26 100644
--- a/stew/assign2.nim
+++ b/stew/assign2.nim
@@ -24,17 +24,30 @@ func assign*[T](tgt: var openArray[T], src: openArray[T]) =
     raiseAssert "Target and source lengths don't match: " &
       $tgt.len & " vs " & $src.len
 
-  assignImpl(tgt, src)
+  when nimvm:
+    for i in 0..<tgt.len:
+      tgt[i] = src[i]
+  else:
+    assignImpl(tgt, src)
 
 func assign*[T](tgt: var seq[T], src: openArray[T]) =
   mixin assign
 
   tgt.setLen(src.len)
-  assignImpl(tgt.toOpenArray(0, tgt.high), src)
+
+  when nimvm:
+    for i in 0..<tgt.len:
+      tgt[i] = src[i]
+  else:
+    assignImpl(tgt.toOpenArray(0, tgt.high), src)
 
 func assign*(tgt: var string, src: string) =
   tgt.setLen(src.len)
-  assignImpl(tgt.toOpenArrayByte(0, tgt.high), src.toOpenArrayByte(0, tgt.high))
+  when nimvm:
+    for i in 0..<tgt.len:
+      tgt[i] = src[i]
+  else:
+    assignImpl(tgt.toOpenArrayByte(0, tgt.high), src.toOpenArrayByte(0, tgt.high))
 
 macro unsupported(T: typed): untyped =
   error "Assignment of the type " & humaneTypeName(T) & " is not supported"
@@ -44,23 +57,25 @@ func assign*[T](tgt: var T, src: T) =
   # is ridiculously slow. When syncing, the application was spending 50%+ CPU
   # time in it - `assign`, in the same test, doesn't even show in the perf trace
   mixin assign
-
-  when supportsCopyMem(T):
-    when sizeof(src) <= sizeof(int):
-      tgt = src
-    else:
-      moveMem(addr tgt, unsafeAddr src, sizeof(tgt))
-  elif T is object|tuple:
-    for t, s in fields(tgt, src):
-      when supportsCopyMem(type s) and sizeof(s) <= sizeof(int) * 2:
-        t = s # Shortcut
-      else:
-        assign(t, s)
-  elif T is seq:
-    assign(tgt, src.toOpenArray(0, src.high))
-  elif T is ref:
+  when nimvm:
     tgt = src
-  elif compiles(distinctBase(tgt)):
-    assign(distinctBase tgt, distinctBase src)
   else:
-    unsupported T
+    when supportsCopyMem(T):
+      when sizeof(src) <= sizeof(int):
+        tgt = src
+      else:
+        moveMem(addr tgt, unsafeAddr src, sizeof(tgt))
+    elif T is object|tuple:
+      for t, s in fields(tgt, src):
+        when supportsCopyMem(type s) and sizeof(s) <= sizeof(int) * 2:
+          t = s # Shortcut
+        else:
+          assign(t, s)
+    elif T is seq:
+      assign(tgt, src.toOpenArray(0, src.high))
+    elif T is ref:
+      tgt = src
+    elif compiles(distinctBase(tgt)):
+      assign(distinctBase tgt, distinctBase src)
+    else:
+      unsupported T
diff --git a/tests/test_assign2.nim b/tests/test_assign2.nim
index 0c61e2b2..e95df110 100644
--- a/tests/test_assign2.nim
+++ b/tests/test_assign2.nim
@@ -39,3 +39,11 @@ suite "assign2":
     assign(dis, X(55))
 
     check: int(dis) == 55
+
+    const x = block: # Compile-time assignment
+      var a, b: array[2, byte]
+      a[1] = 2
+      assign(a, b)
+      b
+   check x[1] == 2
+