Add math_expression and fix ast2 tests

Update to include more expressions

6 files changed, 212 insertions, 78 deletions

diff --git a/nimterop/ast2.nim b/nimterop/ast2.nim
index 04bb8eb..85b555f 100644
--- a/nimterop/ast2.nim
+++ b/nimterop/ast2.nim
@@ -1,10 +1,10 @@
 import macros, os, sequtils, sets, strformat, strutils, tables, times
 
-import compiler/[ast, idents, lineinfos, modulegraphs, msgs, options, parser, renderer]
+import compiler/[ast, idents, lineinfos, modulegraphs, msgs, options, renderer]
 
 import "."/treesitter/api
 
-import "."/[globals, getters, exprparser]
+import "."/[globals, getters, exprparser, utils]
 
 proc getPtrType*(str: string): string =
   result = case str:
@@ -17,21 +17,6 @@ proc getPtrType*(str: string): string =
     else:
       str
 
-proc handleError*(conf: ConfigRef, info: TLineInfo, msg: TMsgKind, arg: string) =
-  # Raise exception in parseString() instead of exiting for errors
-  if msg < warnMin:
-    raise newException(Exception, msgKindToString(msg))
-
-proc parseString(gState: State, str: string): PNode =
-  # Parse a string into Nim AST - use custom error handler that raises
-  # an exception rather than exiting on failure
-  try:
-    result = parseString(
-      str, gState.identCache, gState.config, errorHandler = handleError
-    )
-  except:
-    decho getCurrentExceptionMsg()
-
 proc getLit*(nimState: NimState, str: string, expression = false): PNode =
   result = nimState.codeToNode(str)
 
@@ -121,7 +106,7 @@ proc newConstDef(gState: State, node: TSNode, fname = "", fval = ""): PNode =
   if name.Bl:
     # Name skipped or overridden since blank
     result = nimState.getOverrideOrSkip(node, origname, nskConst)
-  elif valident.kind != nkNilLit:
+  elif valident.kind != nkNone:
     if nimState.addNewIdentifer(name):
       # const X* = Y
       #
@@ -179,7 +164,7 @@ proc addPragma(gState: State, node: TSNode, pragma: PNode, name: string, value:
   if value.isNil:
     pragma.add pident
   else:
-    let
+    var
       colExpr = newNode(nkExprColonExpr)
     colExpr.add pident
     colExpr.add value
@@ -1501,6 +1486,7 @@ proc addProc(gState: State, node, rnode: TSNode) =
       # Parameter list
       plist = node.anyChildInTree("parameter_list")
 
+    var
       procDef = newNode(nkProcDef)
 
     # proc X(a1: Y, a2: Z): P {.pragma.}
diff --git a/nimterop/exprparser.nim b/nimterop/exprparser.nim
index 1f72a9a..505453c 100644
--- a/nimterop/exprparser.nim
+++ b/nimterop/exprparser.nim
@@ -6,7 +6,7 @@ import compiler/[ast, renderer]
 
 import "."/treesitter/[api, c, cpp]
 
-import "."/[globals, getters]
+import "."/[globals, getters, utils]
 
 type
   ExprParser* = ref object
@@ -21,15 +21,16 @@ proc newExprParser*(state: NimState, code: string): ExprParser =
 template techo(msg: varargs[string, `$`]) =
   if exprParser.state.gState.debug:
     let nimState {.inject.} = exprParser.state
-    necho "# " & join(msg, "").replace("\n", "\n# ")
+    necho join(msg, "").getCommented
 
-template val*(node: TSNode): string =
+template val(node: TSNode): string =
   exprParser.code.getNodeVal(node)
 
-proc mode*(exprParser: ExprParser): string =
+proc mode(exprParser: ExprParser): string =
   exprParser.state.gState.mode
 
 template withCodeAst(exprParser: ExprParser, body: untyped): untyped =
+  ## A simple template to inject the TSNode into a body of code
   var parser = tsParserNew()
   defer:
     parser.tsParserDelete()
@@ -51,9 +52,9 @@ template withCodeAst(exprParser: ExprParser, body: untyped): untyped =
   defer:
     tree.tsTreeDelete()
 
-
 proc getNumNode(number, suffix: string): PNode {.inline.} =
-  result = newNode(nkNilLit)
+  ## Convert a C number to a Nim number PNode
+  result = newNode(nkNone)
   if number.contains("."):
     let floatSuffix = number[result.len-1]
     try:
@@ -84,6 +85,8 @@ proc getNumNode(number, suffix: string): PNode {.inline.} =
   else:
     result = newNode(nkIntLit)
 
+  # I realize these regex are wasteful on performance, but
+  # couldn't come up with a better idea.
   if number.contains(re"0[xX]"):
     result.intVal = parseHexInt(number)
     result.flags = {nfBase16}
@@ -97,11 +100,11 @@ proc getNumNode(number, suffix: string): PNode {.inline.} =
     result.intVal = parseInt(number)
 
 proc processNumberLiteral*(exprParser: ExprParser, node: TSNode): PNode =
-  result = newNode(nkNilLit)
+  result = newNode(nkNone)
   let nodeVal = node.val
 
   var match: RegexMatch
-  const reg = re"(\-)?(0\d+|0[xX][0-9a-fA-F]+|0[bB][01]+|\d+\.?\d*[fFlL]?|\d*\.?\d+[fFlL]?|\d+)([ulUL]*)"
+  const reg = re"(\-)?(0\d+|0[xX][0-9a-fA-F]+|0[bB][01]+|\d+|\d+\.?\d*[fFlL]?|\d*\.?\d+[fFlL]?)([ulUL]*)"
   let found = nodeVal.find(reg, match)
   if found:
     let
@@ -111,7 +114,7 @@ proc processNumberLiteral*(exprParser: ExprParser, node: TSNode): PNode =
 
     result = getNumNode(number, suffix)
 
-    if result.kind != nkNilLit and prefix == "-":
+    if result.kind != nkNone and prefix == "-":
       result = nkPrefix.newTree(
         exprParser.state.getIdent("-"),
         result
@@ -127,9 +130,9 @@ proc processStringLiteral*(exprParser: ExprParser, node: TSNode): PNode =
   let nodeVal = node.val
   result = newStrNode(nkStrLit, nodeVal[1 ..< nodeVal.len - 1])
 
-proc processTSNode*(exprParser: ExprParser, node: TSNode): PNode
+proc processTSNode*(exprParser: ExprParser, node: TSNode, typeofNode: PNode = nil): PNode
 
-proc processShiftExpression*(exprParser: ExprParser, node: TSNode): PNode =
+proc processShiftExpression*(exprParser: ExprParser, node: TSNode, typeofNode: PNode = nil): PNode =
   result = newNode(nkInfix)
   let
     left = node[0]
@@ -144,25 +147,29 @@ proc processShiftExpression*(exprParser: ExprParser, node: TSNode): PNode =
   else:
     raise newException(ExprParseError, &"Unsupported shift symbol \"{shiftSym}\"")
 
-  let
-    leftNode = exprParser.processTSNode(left)
-    rightNode = exprParser.processTSNode(right)
+  let leftNode = exprParser.processTSNode(left, typeofNode)
+
+  var tnode = typeofNode
+  if tnode.isNil:
+    tnode = leftNode
+
+  let rightNode = exprParser.processTSNode(right, tnode)
 
   result.add leftNode
   result.add nkCast.newTree(
     nkCall.newTree(
       exprParser.state.getIdent("typeof"),
-      leftNode
+      tnode
     ),
     rightNode
   )
 
-proc processParenthesizedExpr*(exprParser: ExprParser, node: TSNode): PNode =
+proc processParenthesizedExpr*(exprParser: ExprParser, node: TSNode, typeofNode: PNode = nil): PNode =
   result = newNode(nkPar)
   for i in 0 ..< node.len():
-    result.add(exprParser.processTSNode(node[i]))
+    result.add(exprParser.processTSNode(node[i], typeofNode))
 
-proc processLogicalExpression*(exprParser: ExprParser, node: TSNode): PNode =
+proc processLogicalExpression*(exprParser: ExprParser, node: TSNode, typeofNode: PNode = nil): PNode =
   result = newNode(nkPar)
   let child = node[0]
   var nimSym = ""
@@ -179,14 +186,96 @@ proc processLogicalExpression*(exprParser: ExprParser, node: TSNode): PNode =
   techo "LOG CHILD: ", child.val, ", nim: ", nimSym
   result.add nkPrefix.newTree(
     exprParser.state.getIdent(nimSym),
-    exprParser.processTSNode(child)
+    exprParser.processTSNode(child, typeofNode)
   )
 
-proc processBitwiseExpression*(exprParser: ExprParser, node: TSNode): PNode =
+proc processMathExpression(exprParser: ExprParser, node: TSNode, typeofNode: PNode = nil): PNode =
+  if node.len > 1:
+    # Node has left and right children ie: (2 + 7)
+    var
+      res = newNode(nkInfix)
+    let
+      left = node[0]
+      right = node[1]
+
+    let mathSym = exprParser.code[left.tsNodeEndByte() ..< right.tsNodeStartByte()].strip()
+    techo "MATH SYM: ", mathSym
+
+    res.add exprParser.state.getIdent(mathSym)
+    let leftNode = exprParser.processTSNode(left, typeofNode)
+
+    var tnode = typeofNode
+    if tnode.isNil:
+      tnode = leftNode
+
+    let rightNode = exprParser.processTSNode(right, tnode)
+
+    res.add leftNode
+    # res.add rightNode
+    res.add nkCast.newTree(
+      nkCall.newTree(
+        exprParser.state.getIdent("typeof"),
+        tnode
+      ),
+      rightNode
+    )
+
+    # Make sure the statement is of the same type as the left
+    # hand argument, since some expressions return a differing
+    # type than the input types (2/3 == float)
+    result = nkCall.newTree(
+      nkCall.newTree(
+        exprParser.state.getIdent("typeof"),
+        tnode
+      ),
+      res
+    )
+
+  elif node.len() == 1:
+    # Node has only one child, ie -(20 + 7)
+    result = newNode(nkPar)
+    let child = node[0]
+    var nimSym = ""
+
+    let unarySym = exprParser.code[node.tsNodeStartByte() ..< child.tsNodeStartByte()].strip()
+    techo "MATH SYM: ", unarySym
+
+    case unarySym
+    of "+":
+      nimSym = "+"
+    of "-":
+      # Special case. The minus symbol must be in front of an integer,
+      # so we have to make a gental cast here to coerce it to one.
+      # Might be bad because we are overwriting the type
+      # There's probably a better way of doing this
+      result.add nkPrefix.newTree(
+        exprParser.state.getIdent(unarySym),
+        nkPar.newTree(
+          nkCall.newTree(
+            exprParser.state.getIdent("int64"),
+            exprParser.processTSNode(child, typeofNode)
+          )
+        )
+      )
+      return
+    else:
+      raise newException(ExprParseError, &"Unsupported unary symbol \"{unarySym}\"")
+
+    result.add nkPrefix.newTree(
+      exprParser.state.getIdent(nimSym),
+      exprParser.processTSNode(child, typeofNode)
+    )
+  else:
+    raise newException(ExprParseError, &"Invalid bitwise_expression \"{node.val}\"")
+
+proc processBitwiseExpression(exprParser: ExprParser, node: TSNode, typeofNode: PNode = nil): PNode =
   if node.len() > 1:
     result = newNode(nkInfix)
-    let left = node[0]
-    let right = node[1]
+
+    let
+      left = node[0]
+      right = node[1]
+
     var nimSym = ""
 
     var binarySym = exprParser.code[left.tsNodeEndByte() ..< right.tsNodeStartByte()].strip()
@@ -203,15 +292,19 @@ proc processBitwiseExpression*(exprParser: ExprParser, node: TSNode): PNode =
       raise newException(ExprParseError, &"Unsupported binary symbol \"{binarySym}\"")
 
     result.add exprParser.state.getIdent(nimSym)
-    let
-      leftNode = exprParser.processTSNode(left)
-      rightNode = exprParser.processTSNode(right)
+    let leftNode = exprParser.processTSNode(left, typeofNode)
+
+    var tnode = typeofNode
+    if tnode.isNil:
+      tnode = leftNode
+
+    let rightNode = exprParser.processTSNode(right, tnode)
 
     result.add leftNode
-    result.add nkCast.newTree(
+    result.add nkCall.newTree(
       nkCall.newTree(
         exprParser.state.getIdent("typeof"),
-        leftNode
+        tnode
       ),
       rightNode
     )
@@ -221,24 +314,26 @@ proc processBitwiseExpression*(exprParser: ExprParser, node: TSNode): PNode =
     let child = node[0]
     var nimSym = ""
 
-    var binarySym = exprParser.code[node.tsNodeStartByte() ..< child.tsNodeStartByte()].strip()
-    techo "BIN SYM: ", binarySym
+    var unarySym = exprParser.code[node.tsNodeStartByte() ..< child.tsNodeStartByte()].strip()
+    techo "BIN SYM: ", unarySym
 
-    case binarySym
+    case unarySym
     of "~":
       nimSym = "not"
     else:
-      raise newException(ExprParseError, &"Unsupported unary symbol \"{binarySym}\"")
+      raise newException(ExprParseError, &"Unsupported unary symbol \"{unarySym}\"")
 
     result.add nkPrefix.newTree(
       exprParser.state.getIdent(nimSym),
-      exprParser.processTSNode(child)
+      exprParser.processTSNode(child, typeofNode)
     )
   else:
     raise newException(ExprParseError, &"Invalid bitwise_expression \"{node.val}\"")
 
-proc processTSNode*(exprParser: ExprParser, node: TSNode): PNode =
-  result = newNode(nkNilLit)
+proc processTSNode(exprParser: ExprParser, node: TSNode, typeofNode: PNode = nil): PNode =
+  ## Handle all of the types of expressions here. This proc gets called recursively
+  ## in the processX procs and will drill down to sub nodes.
+  result = newNode(nkNone)
   let nodeName = node.getName()
   techo "NODE: ", nodeName, ", VAL: ", node.val
   case nodeName
@@ -251,36 +346,47 @@ proc processTSNode*(exprParser: ExprParser, node: TSNode): PNode =
   of "expression_statement", "ERROR", "translation_unit":
     # This may be wrong. What can be in an expression?
     if node.len > 0:
-      result = exprParser.processTSNode(node[0])
+      result = exprParser.processTSNode(node[0], typeofNode)
     else:
       raise newException(ExprParseError, &"Node type \"{nodeName}\" has no children")
-
   of "parenthesized_expression":
-    result = exprParser.processParenthesizedExpr(node)
+    result = exprParser.processParenthesizedExpr(node, typeofNode)
   of "bitwise_expression":
-    result = exprParser.processBitwiseExpression(node)
+    result = exprParser.processBitwiseExpression(node, typeofNode)
+  of "math_expression":
+    result = exprParser.processMathExpression(node, typeofNode)
   of "shift_expression":
-    result = exprParser.processShiftExpression(node)
+    result = exprParser.processShiftExpression(node, typeofNode)
   of "logical_expression":
-    result = exprParser.processLogicalExpression(node)
+    result = exprParser.processLogicalExpression(node, typeofNode)
+  # Why are these node types named true/false?
+  of "true", "false":
+    result = exprParser.state.parseString(node.val)
   of "identifier":
     var ident = node.val
     if ident != "_":
+      # Process the identifier through cPlugin
       ident = exprParser.state.getIdentifier(ident, nskConst)
-    result = exprParser.state.getIdent(ident)
+      techo ident
+    if ident != "":
+      result = exprParser.state.getIdent(ident)
+    if result.kind == nkNone:
+      raise newException(ExprParseError, &"Could not get identifier \"{ident}\"")
   else:
     raise newException(ExprParseError, &"Unsupported node type \"{nodeName}\" for node \"{node.val}\"")
 
-  techo "NODERES: ", result
+  techo "NODE RES: ", result
 
 proc codeToNode*(state: NimState, code: string): PNode =
-  let exprParser = newExprParser(state, code)
+  ## Convert the C string to a nim PNode tree
+  result = newNode(nkNone)
   try:
+    let exprParser = newExprParser(state, code)
     withCodeAst(exprParser):
       result = exprParser.processTSNode(root)
   except ExprParseError as e:
-    techo e.msg
-    result = newNode(nkNilLit)
+    echo e.msg.getCommented
+    result = newNode(nkNone)
   except Exception as e:
-    techo e.msg
-    result = newNode(nkNilLit)
\ No newline at end of file
+    echo e.msg.getCommented
+    result = newNode(nkNone)
\ No newline at end of file
diff --git a/nimterop/getters.nim b/nimterop/getters.nim
index b96ae70..f43f199 100644
--- a/nimterop/getters.nim
+++ b/nimterop/getters.nim
@@ -399,7 +399,7 @@ proc printLisp*(code: var string, root: TSNode): string =
     else:
       break
 
-    if node.tsNodeNamedChildCount() != 0:
+    if node.len() != 0:
       result &= "\n"
       nextnode = node.tsNodeNamedChild(0)
       depth += 1
@@ -459,15 +459,19 @@ proc printTree*(gState: State, pnode: PNode, offset = ""): string =
     if offset.len == 0:
       result &= "\n"
 
-proc printDebug*(gState: State, node: TSNode) =
-  if gState.debug:
-    gecho ("Input => " & gState.getNodeVal(node)).getCommented() & "\n" &
-          gState.printLisp(node).getCommented()
-
-proc printDebug*(gState: State, pnode: PNode) =
-  if gState.debug:
-    gecho ("Output => " & $pnode).getCommented() & "\n" &
-          gState.printTree(pnode)
+proc printDebug*(nimState: NimState, node: TSNode) =
+  discard
+  # This causes random segfaults for some reason on macOS Catalina
+  if nimState.gState.debug:
+    necho ("Input => " & nimState.getNodeVal(node)).getCommented()
+    necho nimState.gState.printLisp(node).getCommented()
+
+proc printDebug*(nimState: NimState, pnode: PNode) =
+  discard
+  # This causes random segfaults for some reason on macOS Catalina
+  if nimState.gState.debug and pnode.kind != nkNone:
+    necho ("Output => " & $pnode).getCommented()
+    necho nimState.printTree(pnode)
 
 # Compiler shortcuts
 
diff --git a/nimterop/globals.nim b/nimterop/globals.nim
index f159124..b964da9 100644
--- a/nimterop/globals.nim
+++ b/nimterop/globals.nim
@@ -113,7 +113,7 @@ when not declared(CIMPORT):
   export gAtoms, gExpressions, gEnumVals, Kind, Ast, AstTable, State, nBl, Bl
 
   # Redirect output to file when required
-  template gecho*(args: string) {.dirty.} =
+  template gecho*(args: string) =
     if gState.outputHandle.isNil:
       echo args
     else:
diff --git a/tests/include/tast2.h b/tests/include/tast2.h
index bdf8823..7e15ac0 100644
--- a/tests/include/tast2.h
+++ b/tests/include/tast2.h
@@ -8,6 +8,24 @@ extern "C" {
 #define D "hello"
 #define E 'c'
 
+#define UEXPR (1234u << 1)
+#define ULEXPR (1234ul << 2)
+#define ULLEXPR (1234ull << 3)
+#define LEXPR (1234l << 4)
+#define LLEXPR (1234ll << 5)
+
+#define SHL1 (1u << 1)
+#define SHL2 (1u << 2)
+#define SHL3 (1u << 3)
+#define COERCE 645635634896ull + -35436
+#define COERCE2 645635634896 + -35436
+#define BINEXPR ~(-(1u << !-1)) ^ (10 >> 1)
+#define BOOL true
+#define MATHEXPR (1 + 2/3*20 - 100)
+#define ANDEXPR (100 & 11000)
+
+#define ALLSHL (SHL1 | SHL2 | SHL3)
+
 struct A0;
 struct A1 {};
 typedef struct A2;
diff --git a/tests/tast2.nim b/tests/tast2.nim
index e13c4ac..d65e34a 100644
--- a/tests/tast2.nim
+++ b/tests/tast2.nim
@@ -105,6 +105,26 @@ assert C == 0x10
 assert D == "hello"
 assert E == 'c'
 
+assert UEXPR == (1234.uint shl 1)
+assert ULEXPR == (1234.uint32 shl 2)
+assert ULLEXPR == (1234.uint64 shl 3)
+assert LEXPR == (1234.int32 shl 4)
+assert LLEXPR == (1234.int64 shl 5)
+
+assert COERCE == 645635599460'u64
+assert COERCE2 == 645635599460'i64
+
+assert BINEXPR == 5
+assert BOOL == true
+assert MATHEXPR == -99
+assert ANDEXPR == 96
+
+assert SHL1 == (1.uint shl 1)
+assert SHL2 == (1.uint shl 2)
+assert SHL3 == (1.uint shl 3)
+
+assert ALLSHL == (SHL1 or SHL2 or SHL3)
+
 assert A0 is object
 testFields(A0, "f1!cint")
 checkPragmas(A0, pHeaderBy, istype = false)
@@ -271,7 +291,7 @@ var a21p: A21p
 a21p = addr a20
 
 assert A22 is object
-testFields(A22, "f1|f2!ptr ptr cint|array[123 + 132, ptr cint]")
+testFields(A22, "f1|f2!ptr ptr cint|array[typeof(123)(123 + cast[typeof(123)](132)), ptr cint]")
 checkPragmas(A22, pHeaderBy, istype = false)
 var a22: A22
 a22.f1 = addr a15.a2[0]