From 1b143f5e79c940ba7f70e0512f36b5c61a6bc24d Mon Sep 17 00:00:00 2001
From: hlaaftana <10591326+hlaaftana@users.noreply.github.com>
Date: Mon, 22 Nov 2021 12:41:35 +0300
Subject: [PATCH] Accept nnkTypeSection from typedef macro pragmas (#19168)

---
 changelog.md                     | 31 ++++++++++++++-
 compiler/semstmts.nim            | 18 ++++++---
 doc/manual.rst                   |  1 +
 tests/pragmas/ttypedef_macro.nim | 66 ++++++++++++++++++++++++++++++++
 4 files changed, 110 insertions(+), 6 deletions(-)
 create mode 100644 tests/pragmas/ttypedef_macro.nim

diff --git a/changelog.md b/changelog.md
index 1636d125ff1f..8e70513e01e8 100644
--- a/changelog.md
+++ b/changelog.md
@@ -13,7 +13,36 @@
 
 ## Language changes
 
-
+- Pragma macros on type definitions can now return `nnkTypeSection` nodes as well as `nnkTypeDef`,
+  allowing multiple type definitions to be injected in place of the original type definition.
+
+  ```nim
+  import macros
+
+  macro multiply(amount: static int, s: untyped): untyped =
+    let name = $s[0].basename
+    result = newNimNode(nnkTypeSection)
+    for i in 1 .. amount:
+      result.add(newTree(nnkTypeDef, ident(name & $i), s[1], s[2]))
+
+  type
+    Foo = object
+    Bar {.multiply: 3.} = object
+      x, y, z: int
+    Baz = object
+
+  # becomes
+
+  type
+    Foo = object
+    Bar1 = object
+      x, y, z: int
+    Bar2 = object
+      x, y, z: int
+    Bar3 = object
+      x, y, z: int
+    Baz = object
+  ```
 
 ## Compiler changes
 
diff --git a/compiler/semstmts.nim b/compiler/semstmts.nim
index 4607e857aa10..36d76608e39f 100644
--- a/compiler/semstmts.nim
+++ b/compiler/semstmts.nim
@@ -1112,7 +1112,12 @@ proc typeDefLeftSidePass(c: PContext, typeSection: PNode, i: int) =
     if name.kind == nkPragmaExpr:
       let rewritten = applyTypeSectionPragmas(c, name[1], typeDef)
       if rewritten != nil:
-        typeSection[i] = rewritten
+        case rewritten.kind
+        of nkTypeDef:
+          typeSection[i] = rewritten
+        of nkTypeSection:
+          typeSection.sons[i .. i] = rewritten.sons
+        else: illFormedAst(rewritten, c.config)
         typeDefLeftSidePass(c, typeSection, i)
         return
       pragma(c, s, name[1], typePragmas)
@@ -1143,16 +1148,19 @@ proc typeDefLeftSidePass(c: PContext, typeSection: PNode, i: int) =
 proc typeSectionLeftSidePass(c: PContext, n: PNode) =
   # process the symbols on the left side for the whole type section, before
   # we even look at the type definitions on the right
-  for i in 0..<n.len:
+  var i = 0
+  while i < n.len: # n may grow due to type pragma macros
     var a = n[i]
     when defined(nimsuggest):
       if c.config.cmd == cmdIdeTools:
         inc c.inTypeContext
         suggestStmt(c, a)
         dec c.inTypeContext
-    if a.kind == nkCommentStmt: continue
-    if a.kind != nkTypeDef: illFormedAst(a, c.config)
-    typeDefLeftSidePass(c, n, i)
+    case a.kind
+    of nkCommentStmt: discard
+    of nkTypeDef: typeDefLeftSidePass(c, n, i)
+    else: illFormedAst(a, c.config)
+    inc i
 
 proc checkCovariantParamsUsages(c: PContext; genericType: PType) =
   var body = genericType[^1]
diff --git a/doc/manual.rst b/doc/manual.rst
index 328035b6d4f1..a4701c85bddb 100644
--- a/doc/manual.rst
+++ b/doc/manual.rst
@@ -7755,6 +7755,7 @@ This is translated to:
 This is translated to a call to the `schema` macro with a `nnkTypeDef`
 AST node capturing both the left-hand side and right-hand side of the
 definition. The macro can return a potentially modified `nnkTypeDef` tree
+or multiple `nnkTypeDef` trees contained in a `nnkTypeSection` node
 which will replace the original row in the type section.
 
 When multiple macro pragmas are applied to the same definition, the
diff --git a/tests/pragmas/ttypedef_macro.nim b/tests/pragmas/ttypedef_macro.nim
new file mode 100644
index 000000000000..dd4c8775700f
--- /dev/null
+++ b/tests/pragmas/ttypedef_macro.nim
@@ -0,0 +1,66 @@
+import macros
+
+macro makeref(s): untyped =
+  expectKind s, nnkTypeDef
+  result = newTree(nnkTypeDef, s[0], s[1], newTree(nnkRefTy, s[2]))
+
+type
+  Obj {.makeref.} = object
+    a: int
+
+doAssert Obj is ref
+doAssert Obj(a: 3)[].a == 3
+
+macro multiply(amount: static int, s): untyped =
+  let name = $s[0].basename
+  result = newNimNode(nnkTypeSection)
+  for i in 1 .. amount:
+    result.add(newTree(nnkTypeDef, ident(name & $i), s[1], s[2]))
+
+type
+  Foo = object
+  Bar {.multiply: 2.} = object
+    x, y, z: int
+  Baz = object
+
+let bar1 = Bar1(x: 1, y: 2, z: 3)
+let bar2 = Bar2(x: bar1.x, y: bar1.y, z: bar1.z)
+doAssert Bar1 isnot Bar2
+doAssert not declared(Bar)
+doAssert not declared(Bar3)
+
+# https://github.com/nim-lang/RFCs/issues/219
+
+macro inferKind(td): untyped =
+  let name = $td[0].basename
+  var rhs = td[2]
+  while rhs.kind in {nnkPtrTy, nnkRefTy}: rhs = rhs[0]
+  if rhs.kind != nnkObjectTy:
+    result = td
+  else:
+    for n in rhs[^1]:
+      if n.kind == nnkRecCase and n[0][^2].eqIdent"_":
+        let kindTypeName = ident(name & "Kind")
+        let en = newTree(nnkEnumTy, newEmptyNode())
+        for i in 1 ..< n.len:
+          let branch = n[i]
+          if branch.kind == nnkOfBranch:
+            for j in 0 ..< branch.len - 1:
+              en.add(branch[j])
+        n[0][^2] = kindTypeName
+        return newTree(nnkTypeSection,
+          newTree(nnkTypeDef, kindTypeName, newEmptyNode(), en),
+          td)
+
+type Node {.inferKind.} = ref object
+  case kind: _
+  of opValue: value: int
+  of opAdd, opSub, opMul, opCall: kids: seq[Node]
+
+doAssert opValue is NodeKind
+let node = Node(kind: opMul, kids: @[
+  Node(kind: opValue, value: 3),
+  Node(kind: opValue, value: 5)
+])
+doAssert node.kind == opMul
+doAssert node.kids[0].value * node.kids[1].value == 15