From 4ffd3e1f59c2f12ae57456c6fccf477bc203da78 Mon Sep 17 00:00:00 2001
From: Zahary Karadjov <zahary@gmail.com>
Date: Tue, 26 May 2020 18:42:33 +0300
Subject: [PATCH] Support type hierarchies in recordFields

---
 stew/shims/macros.nim | 72 ++++++++++++++++++++++++++++---------------
 tests/all_tests.nim   |  1 +
 tests/test_macros.nim | 38 +++++++++++++++++++++++
 3 files changed, 87 insertions(+), 24 deletions(-)
 create mode 100644 tests/test_macros.nim

diff --git a/stew/shims/macros.nim b/stew/shims/macros.nim
index 4f143d4..8fc4541 100644
--- a/stew/shims/macros.nim
+++ b/stew/shims/macros.nim
@@ -74,45 +74,51 @@ func isTuple*(t: NimNode): bool =
 macro isTuple*(T: type): untyped =
   newLit(isTuple(getType(T)[1]))
 
+proc skipRef*(T: NimNode): NimNode =
+  result = T
+  if T.kind == nnkBracketExpr:
+    if eqIdent(T[0], bindSym"ref"):
+      result = T[1]
+
 template readPragma*(field: FieldDescription, pragmaName: static string): NimNode =
   let p = findPragma(field.pragmas, bindSym(pragmaName))
   if p != nil and p.len == 2: p[1] else: p
 
-proc recordFieldsAux(result: var seq[FieldDescription],
-                     n: NimNode,
-                     parentCaseField: NimNode = nil,
-                     parentCaseBranch: NimNode = nil,
-                     isDiscriminator = false) =
+proc collectFieldsFromRecList(result: var seq[FieldDescription],
+                              n: NimNode,
+                              parentCaseField: NimNode = nil,
+                              parentCaseBranch: NimNode = nil,
+                              isDiscriminator = false) =
   case n.kind
   of nnkRecList:
     for entry in n:
-      recordFieldsAux result, entry,
-                      parentCaseField, parentCaseBranch
+      collectFieldsFromRecList result, entry,
+                               parentCaseField, parentCaseBranch
   of nnkRecWhen:
     for branch in n:
       case branch.kind:
       of nnkElifBranch:
-        recordFieldsAux result, branch[1],
-                        parentCaseField, parentCaseBranch
+        collectFieldsFromRecList result, branch[1],
+                                 parentCaseField, parentCaseBranch
       of nnkElse:
-        recordFieldsAux result, branch[0],
-                        parentCaseField, parentCaseBranch
+        collectFieldsFromRecList result, branch[0],
+                                 parentCaseField, parentCaseBranch
       else:
         doAssert false
 
   of nnkRecCase:
-    recordFieldsAux result, n[0],
-                    parentCaseField,
-                    parentCaseBranch,
-                    isDiscriminator = true
+    collectFieldsFromRecList result, n[0],
+                             parentCaseField,
+                             parentCaseBranch,
+                             isDiscriminator = true
 
     for i in 1 ..< n.len:
       let branch = n[i]
       case branch.kind
       of nnkOfBranch:
-        recordFieldsAux result, branch[^1], n[0], branch
+        collectFieldsFromRecList result, branch[^1], n[0], branch
       of nnkElse:
-        recordFieldsAux result, branch[0], n[0], branch
+        collectFieldsFromRecList result, branch[0], n[0], branch
       else:
         doAssert false
 
@@ -150,19 +156,37 @@ proc recordFieldsAux(result: var seq[FieldDescription],
   else:
     doAssert false, "Unexpected nodes in recordFields:\n" & n.treeRepr
 
+proc collectFieldsInHierarchy(result: var seq[FieldDescription],
+                              objectType: NimNode) =
+  var objectType = objectType
+
+  if objectType.kind == nnkRefTy:
+    objectType = objectType[0]
+
+  objectType.expectKind nnkObjectTy
+
+  var baseType = objectType[1]
+  if baseType.kind != nnkEmpty:
+    baseType.expectKind nnkOfInherit
+    baseType = baseType[0]
+    baseType.expectKind nnkSym
+    baseType = getImpl(baseType)
+    baseType.expectKind nnkTypeDef
+    baseType = baseType[2]
+    baseType.expectKind nnkObjectTy
+    collectFieldsInHierarchy result, baseType
+
+  let recList = objectType[2]
+  collectFieldsFromRecList result, recList
+
 proc recordFields*(typeImpl: NimNode): seq[FieldDescription] =
-  # TODO: This doesn't support inheritance yet
   if typeImpl.isTuple:
     for i in 1 ..< typeImpl.len:
       result.add FieldDescription(typ: typeImpl[i], name: ident("Field" & $(i - 1)))
     return
 
-  var objectType = typeImpl[2]
-  if objectType.kind == nnkRefTy:
-    objectType = objectType[0]
-
-  let recList = objectType[2]
-  recordFieldsAux result, recList
+  typeImpl.expectKind nnkTypeDef
+  collectFieldsInHierarchy(result, typeImpl[2])
 
 macro field*(obj: typed, fieldName: static string): untyped =
   newDotExpr(obj, ident fieldName)
diff --git a/tests/all_tests.nim b/tests/all_tests.nim
index 46a6d2d..1ff5867 100644
--- a/tests/all_tests.nim
+++ b/tests/all_tests.nim
@@ -16,6 +16,7 @@ import
   test_bitseqs,
   test_byteutils,
   test_endians2,
+  test_macros,
   test_objects,
   test_ptrops,
   test_results,
diff --git a/tests/test_macros.nim b/tests/test_macros.nim
new file mode 100644
index 0000000..ebb31d2
--- /dev/null
+++ b/tests/test_macros.nim
@@ -0,0 +1,38 @@
+import
+  ../stew/shims/macros
+
+type
+  FieldKind = enum
+    KindA
+    KindB
+
+  BaseType = object of RootObj
+    baseField: int
+    case baseCaseField: FieldKind
+    of KindA:
+      baseA: int
+    of KindB:
+      discard
+
+  DerivedType = ref object of BaseType
+    derivedField: int
+
+macro getFieldsLists(T: type): untyped =
+  result = newTree(nnkBracket)
+
+  var resolvedType = skipRef getType(T)[1]
+  doAssert resolvedType.kind == nnkSym
+  var objectType = getImpl(resolvedType)
+  doAssert objectType.kind == nnkTypeDef
+
+  for f in recordFields(objectType):
+    result.add newLit($f.name)
+
+static:
+  doAssert getFieldsLists(DerivedType) == [
+    "baseField",
+    "baseCaseField",
+    "baseA",
+    "derivedField"
+  ]
+