Initial ast2 pragma support

5 files changed, 107 insertions, 60 deletions

diff --git a/nimterop/ast.nim b/nimterop/ast.nim
index e28241f..06b0504 100644
--- a/nimterop/ast.nim
+++ b/nimterop/ast.nim
@@ -4,6 +4,42 @@ import regex
 
 import "."/[getters, globals, treesitter/api]
 
+proc getHeaderPragma*(nimState: NimState): string =
+  result =
+    if nimState.includeHeader():
+      &", header: {nimState.currentHeader}"
+    else:
+      ""
+
+proc getDynlib*(nimState: NimState): string =
+  result =
+    if nimState.gState.dynlib.nBl:
+      &", dynlib: {nimState.gState.dynlib}"
+    else:
+      ""
+
+proc getImportC*(nimState: NimState, origName, nimName: string): string =
+  if nimName != origName:
+    result = &"importc: \"{origName}\"{nimState.getHeaderPragma()}"
+  else:
+    result = nimState.impShort
+
+proc getPragma*(nimState: NimState, pragmas: varargs[string]): string =
+  result = ""
+  for pragma in pragmas.items():
+    if pragma.nBl:
+      result &= pragma & ", "
+  if result.nBl:
+    result = " {." & result[0 .. ^3] & ".}"
+
+  result = result.replace(nimState.impShort & ", cdecl", nimState.impShort & "C")
+
+  let
+    dy = nimState.getDynlib()
+
+  if ", cdecl" in result and dy.nBl:
+    result = result.replace(".}", dy & ".}")
+
 proc saveNodeData(node: TSNode, nimState: NimState): bool =
   let name = $node.tsNodeType()
 
@@ -178,7 +214,7 @@ proc printNim*(gState: State, fullpath: string, root: TSNode, astTable: AstTable
   nimState.impShort = nimState.currentHeader.replace("header", "imp")
   nimState.sourceFile = fullpath
 
-  if nimState.gState.dynlib.Bl and nimState.gState.includeHeader:
+  if nimState.includeHeader():
     nimState.constStr &= &"\n  {nimState.currentHeader} {{.used.}} = \"{fp}\""
 
   root.searchAst(astTable, nimState)
diff --git a/nimterop/ast2.nim b/nimterop/ast2.nim
index 536ef9b..8804c53 100644
--- a/nimterop/ast2.nim
+++ b/nimterop/ast2.nim
@@ -107,7 +107,23 @@ proc addConst(nimState: NimState, node: TSNode) =
 
       nimState.printDebug(constDef)
 
-proc newPragma(nimState: NimState, node: TSNode, pragmas: OrderedTable[string, string]): PNode =
+proc addPragma(nimState: NimState, node: TSNode, pragma: PNode, pragmas: OrderedTable[string, PNode]) =
+  # Add pragmas to an existing nkPragma tree
+  for name, value in pragmas.pairs:
+    let
+      (_, pinfo) = nimState.getNameInfo(node.getAtom(), nskUnknown)
+      pident = nimState.getIdent(name, pinfo, exported = false)
+    
+    if value.isNil:
+      pragma.add pident
+    else:
+      let
+        colExpr = newNode(nkExprColonExpr)
+      colExpr.add pident
+      colExpr.add value
+      pragma.add colExpr
+
+proc newPragma(nimState: NimState, node: TSNode, pragmas: OrderedTable[string, PNode]): PNode =
   # Create nkPragma tree for name:value
   #
   # {.name1, name2: value2.}
@@ -120,22 +136,9 @@ proc newPragma(nimState: NimState, node: TSNode, pragmas: OrderedTable[string, s
   #  )
   # )
   result = newNode(nkPragma)
-  for name, value in pragmas.pairs:
-    let
-      (_, pinfo) = nimState.getNameInfo(node.getAtom(), nskUnknown)
-      pident = nimState.getIdent(name, pinfo, exported = false)
-    
-    if value.len == 0:
-      result.add pident
-    else:
-      let
-        colExpr = newNode(nkExprColonExpr)
-        pvalue = newStrNode(nkStrLit, value)
-      colExpr.add pident
-      colExpr.add pvalue
-      result.add colExpr
+  nimState.addPragma(node, result, pragmas)
 
-proc newPragmaExpr(nimState: NimState, node: TSNode, ident: PNode, pragmas: OrderedTable[string, string]): PNode =
+proc newPragmaExpr(nimState: NimState, node: TSNode, ident: PNode, pragmas: OrderedTable[string, PNode]): PNode =
   # Create nkPragmaExpr tree
   #
   # nkPragmaExpr(
@@ -169,7 +172,9 @@ proc newTypeIdent(nimState: NimState, node: TSNode, override = "", union = false
         nimState.getIdent(name, info)
     prident =
       if union:
-        nimState.newPragmaExpr(node, ident, {"union": ""}.toOrderedTable())
+        var
+          empty: PNode
+        nimState.newPragmaExpr(node, ident, {"union": empty}.toOrderedTable())
       else:
         ident
 
@@ -1008,6 +1013,39 @@ proc searchTree(nimState: NimState, root: TSNode) =
     if node == root:
       break
 
+proc setupPragmas(nimState: NimState, root: TSNode, fullpath: string) =
+  var
+    impPragma = newNode(nkPragma)
+    impCPragma = newNode(nkPragma)
+    empty: PNode
+
+  nimState.addPragma(root, impPragma, {
+    "pragma": nimState.getIdent(nimState.impShort),
+    "importc": empty
+  }.toOrderedTable())
+
+  if nimState.includeHeader():
+    nimState.constSection.add nimState.newConstDef(
+      root, name = nimState.currentHeader, val = fullpath)
+
+    nimState.addPragma(root, impPragma, {
+      "header": newStrNode(nkStrLit, nimState.currentHeader)
+    }.toOrderedTable())
+
+  nimState.addPragma(root, impCPragma, {
+    "pragma": nimState.getIdent(nimState.impShort & "C"),
+    nimState.impShort: empty,
+    "cdecl": empty
+  }.toOrderedTable())
+
+  if nimState.gState.dynlib.nBl:
+    nimState.addPragma(root, impCPragma, {
+      "dynlib": nimState.getIdent(nimState.gState.dynlib)
+    }.toOrderedTable())
+
+  nimState.pragmaSection.add impPragma
+  nimState.pragmaSection.add impCPragma
+
 proc printNimHeader*(gState: State) =
   gecho """# Generated at $1
 # Command line:
@@ -1035,19 +1073,19 @@ proc printNim*(gState: State, fullpath: string, root: TSNode) =
   nimState.config = newConfigRef()
   nimstate.graph = newModuleGraph(nimState.identCache, nimState.config)
 
+  nimState.pragmaSection = newNode(nkStmtList)
   nimState.constSection = newNode(nkConstSection)
   nimState.enumSection = newNode(nkStmtList)
   nimState.procSection = newNode(nkStmtList)
   nimState.typeSection = newNode(nkTypeSection)
 
-  if nimState.gState.dynlib.Bl and nimState.gState.includeHeader:
-    nimState.constSection.add nimState.newConstDef(
-      root, name = nimState.currentHeader, val = fp)
+  nimState.setupPragmas(root, fp)
 
   nimState.searchTree(root)
 
   var
     tree = newNode(nkStmtList)
+  tree.add nimState.pragmaSection
   tree.add nimState.enumSection
   tree.add nimState.constSection
   tree.add nimState.typeSection
diff --git a/nimterop/getters.nim b/nimterop/getters.nim
index 2c7da75..07ed194 100644
--- a/nimterop/getters.nim
+++ b/nimterop/getters.nim
@@ -426,6 +426,9 @@ proc printDebug*(nimState: NimState, pnode: PNode) =
 
 # Compiler shortcuts
 
+proc getDefaultLineInfo*(nimState: NimState): TLineInfo =
+  result = newLineInfo(nimState.config, nimState.sourceFile.AbsoluteFile, -1, -1)
+
 proc getLineInfo*(nimState: NimState, node: TSNode): TLineInfo =
   # Get Nim equivalent line:col info from node
   let
@@ -446,6 +449,9 @@ proc getIdent*(nimState: NimState, name: string, info: TLineInfo, exported = tru
   else:
     result = newIdentNode(ident, info)
 
+proc getIdent*(nimState: NimState, name: string): PNode =
+  nimState.getIdent(name, nimState.getDefaultLineInfo(), exported = false)
+
 proc getNameInfo*(nimState: NimState, node: TSNode, kind: NimSymKind, parent = ""):
   tuple[name: string, info: TLineInfo] =
   # Shortcut to get identifier name and info (node value and line:col)
@@ -648,41 +654,8 @@ proc getSplitComma*(joined: seq[string]): seq[string] =
   for i in joined:
     result = result.concat(i.split(","))
 
-proc getHeaderPragma*(nimState: NimState): string =
-  result =
-    if nimState.gState.dynlib.Bl and nimState.gState.includeHeader:
-      &", header: {nimState.currentHeader}"
-    else:
-      ""
-
-proc getDynlib*(nimState: NimState): string =
-  result =
-    if nimState.gState.dynlib.nBl:
-      &", dynlib: {nimState.gState.dynlib}"
-    else:
-      ""
-
-proc getImportC*(nimState: NimState, origName, nimName: string): string =
-  if nimName != origName:
-    result = &"importc: \"{origName}\"{nimState.getHeaderPragma()}"
-  else:
-    result = nimState.impShort
-
-proc getPragma*(nimState: NimState, pragmas: varargs[string]): string =
-  result = ""
-  for pragma in pragmas.items():
-    if pragma.nBl:
-      result &= pragma & ", "
-  if result.nBl:
-    result = " {." & result[0 .. ^3] & ".}"
-
-  result = result.replace(nimState.impShort & ", cdecl", nimState.impShort & "C")
-
-  let
-    dy = nimState.getDynlib()
-
-  if ", cdecl" in result and dy.nBl:
-    result = result.replace(".}", dy & ".}")
+template includeHeader*(nimState: NimState): bool =
+  nimState.gState.dynlib.Bl and nimState.gState.includeHeader
 
 proc getComments*(nimState: NimState, strip = false): string =
   if not nimState.gState.nocomments and nimState.commentStr.nBl:
diff --git a/nimterop/globals.nim b/nimterop/globals.nim
index ba63a76..2fb543b 100644
--- a/nimterop/globals.nim
+++ b/nimterop/globals.nim
@@ -78,7 +78,7 @@ type
 
     # Nim compiler objects
     when not declared(CIMPORT):
-      constSection*, enumSection*, procSection*, typeSection*: PNode
+      pragmaSection*, constSection*, enumSection*, procSection*, typeSection*: PNode
       identCache*: IdentCache
       config*: ConfigRef
       graph*: ModuleGraph
diff --git a/nimterop/grammar.nim b/nimterop/grammar.nim
index f15c2f8..2e6b585 100644
--- a/nimterop/grammar.nim
+++ b/nimterop/grammar.nim
@@ -2,7 +2,7 @@ import macros, strformat, strutils, tables
 
 import regex
 
-import "."/[getters, globals, lisp, treesitter/api]
+import "."/[ast, getters, globals, lisp, treesitter/api]
 
 type
   Grammar = seq[tuple[grammar: string, call: proc(ast: ref Ast, node: TSNode, nimState: NimState) {.nimcall.}]]
@@ -200,7 +200,7 @@ proc initGrammar(): Grammar =
         nname = nimState.getIdentifier(name, nskType)
         i += 1
 
-      if nimState.gState.dynlib.Bl and nimState.gState.includeHeader:
+      if nimState.includeHeader():
         pragmas.add nimState.getImportC(name, nname)
 
       let
@@ -316,7 +316,7 @@ proc initGrammar(): Grammar =
       else:
         var
           pragmas: seq[string] = @[]
-        if nimState.gState.dynlib.Bl and nimState.gState.includeHeader:
+        if nimState.includeHeader():
           pragmas.add nimState.getImportC(prefix & name, nname)
         pragmas.add "bycopy"
         if union.nBl: