treeform better macro

src/pixie/simd.nim

@@ -1,43 +1,9 @@
-import simd/internal, std/macros, std/tables
+import simd/internal
+
+export internal
 
 const allowSimd* = not defined(pixieNoSimd) and not defined(tcc)
 
-macro hasSimd*(procedure: untyped) =
-  let
-    name = procedure.procName()
-    args = procedure.procArguments()
-    originalBody = procedure[6]
-    nameSse2 = name & "Sse2"
-    nameAvx = name & "Avx"
-    nameAvx2 = name & "Avx2"
-    callAvx = call(ident(nameAvx), args)
-    callAvx2 = call(ident(nameAvx2), args)
-
-  var body = newStmtList()
-
-  when not defined(pixieNoAvx):
-    if nameAvx2 in simdProcs:
-      body.add quote do:
-        if cpuHasAvx2:
-          forceReturn `callAvx2`
-
-    if nameAvx in simdProcs:
-      body.add quote do:
-        if cpuHasAvx:
-          forceReturn `callAvx`
-
-  if nameSse2 in simdProcs:
-    let bodySse2 = simdProcs[nameSse2][6]
-    body.add quote do:
-      `bodySse2`
-  else:
-    body.add quote do:
-      `originalBody`
-
-  procedure[6] = body
-
-  return procedure
-
 when allowSimd and defined(amd64):
   import simd/sse2, simd/avx, simd/avx2
   export sse2, avx, avx2

src/pixie/simd/internal.nim

@@ -2,16 +2,7 @@ import std/macros, std/tables
 
 var simdProcs* {.compiletime.}: Table[string, NimNode]
 
-template forceReturn*(procedure: untyped) =
+proc procName(procedure: NimNode): string =
-  ## Produce `return procedure()` when procedure returns something otherwise
-  ## `procedure(); return` if it procedure returns nothing.
-  when compiles(block: return procedure):
-    return procedure
-  else:
-    procedure
-    return
-
-proc procName*(procedure: NimNode): string =
   ## Given a procedure signature returns only name string.
   let nameNode = procedure[0]
   if nameNode.kind == nnkPostfix:
@@ -19,21 +10,69 @@ proc procName*(procedure: NimNode): string =
   else:
     nameNode.strVal
 
-proc procArguments*(procedure: NimNode): seq[NimNode] =
+proc procArguments(procedure: NimNode): seq[NimNode] =
   ## Given a procedure signature gets the arguments as a list.
   for i, arg in procedure[3]:
     if i > 0:
       for j in 0 ..< arg.len - 2:
         result.add(arg[j])
 
-proc call*(name: NimNode, args: seq[NimNode]): NimNode =
+proc procReturnType(procedure: NimNode): NimNode =
+  ## Given a procedure signature gets the return type.
+  procedure[3][0]
+
+proc callAndReturn(name: NimNode, procedure: NimNode): NimNode =
   ## Produces a procedure call with arguments.
-  result = newNimNode(nnkCall)
+  let
-  result.add(name)
+    retType = procedure.procReturnType()
-  for arg in args:
+    call = newNimNode(nnkCall)
-    result.add(arg)
+  call.add(name)
+  for arg in procedure.procArguments():
+    call.add(arg)
+  if retType.kind == nnkEmpty:
+    result = quote do:
+      `call`
+      return
+  else:
+    result = quote do:
+      return `call`
 
 macro simd*(procedure: untyped) =
   let name = procedure.procName()
   simdProcs[name] = procedure.copy()
   return procedure
+
+macro hasSimd*(procedure: untyped) =
+  let
+    name = procedure.procName()
+    originalBody = procedure[6]
+    nameSse2 = name & "Sse2"
+    nameAvx = name & "Avx"
+    nameAvx2 = name & "Avx2"
+    callAvx = callAndReturn(ident(nameAvx), procedure)
+    callAvx2 = callAndReturn(ident(nameAvx2), procedure)
+
+  var body = newStmtList()
+
+  when not defined(pixieNoAvx):
+    if nameAvx2 in simdProcs:
+      body.add quote do:
+        if cpuHasAvx2:
+          `callAvx2`
+
+    if nameAvx in simdProcs:
+      body.add quote do:
+        if cpuHasAvx2:
+          `callAvx`
+
+  if nameSse2 in simdProcs:
+    let bodySse2 = simdProcs[nameSse2][6]
+    body.add quote do:
+      `bodySse2`
+  else:
+    body.add quote do:
+      `originalBody`
+
+  procedure[6] = body
+
+  return procedure