diff --git a/README.md b/README.md index 4fc5fc1..a6e63e4 100644 --- a/README.md +++ b/README.md @@ -9,7 +9,7 @@ Extensions so far include: This module was modified from the OOP section on the excellent website http://nim-by-example.github.io/ - usage: + Usage: ```nim import extensions/oop @@ -35,3 +35,21 @@ Extensions so far include: ```bash nim c -r tests/test1.nim ``` + +## Enumerate + + This macro is useful for enumerating on any iterator. It will autodetect which iterator you want to use and give you a useful enumeration. + + Usage: + + ```nim + import os + import extensions/enumerate + + for i, line in enumerate(lines, "file.txt"): + echo i, " ", line + + for i, tup in enumerate(walkDir, "."): + echo i, " ", tup + ``` + diff --git a/extensions.nimble b/extensions.nimble index 0f3e7ae..07e3c22 100644 --- a/extensions.nimble +++ b/extensions.nimble @@ -1,6 +1,6 @@ # Package -version = "0.2.0" +version = "0.2.1" author = "Joey Yakimowich-Payne" description = "Nim extensions" license = "MIT" @@ -8,4 +8,4 @@ srcDir = "src" # Dependencies -requires "nim >= 0.18.0" +requires "nim >= 0.18.1" diff --git a/src/extensions/enumerate.nim b/src/extensions/enumerate.nim new file mode 100644 index 0000000..3f64285 --- /dev/null +++ b/src/extensions/enumerate.nim @@ -0,0 +1,75 @@ +import macros +import extensions/nimnode + +proc argsMatch(impl, args: NimNode): bool = + var implArgs = impl["FormalParams"] + result = true + for i, arg in args.pairs: + let implIndex = i + 1 # skip the return type + let symType = arg.getType() + let implType = implArgs[implIndex][1] + result = result and ($symType == $implType) + +macro enumerate*(iter: typed, args: varargs[typed]): untyped = + ## Takes in an i + var iterDef: NimNode + + if len(iter) > 1: + # If there are multiple syms to choose from, + # choose the one with matching args + for c in iter: + let impl = c.getImpl().copy() + if argsMatch(impl, args): + iterDef = impl + else: + iterDef = iter.getImpl().copy() + + # set the symbol to empty because + # we don't need it + iterDef["Sym"] = newEmptyNode() + + # remove all stmts, we only want the def + # for now + iterDef["StmtList"] = newEmptyNode() + + # remove pragma + iterDef["Pragma"] = newEmptyNode() + + var + returnType = iterDef["FormalParams > 0"] + iterCall = newNimNode(nnkCall).add(iter.copy()) + + for child in iterDef[3].children: + if child.kind == nnkIdentDefs: + iterCall.add(child[0].copy()) + + # mod the return type to return both int and original type + iterDef["FormalParams > 0"] = newNimNode(nnkPar).add(ident"int", returnType) + #iterDef[3][0] = newNimNode(nnkPar).add(ident"int", returnType) + + template iterBody(iterCall) = + var i = 0 + for tup in iterCall: + # Since Nim's checker won't allow a yield statement here, + # we need to get creative + replace(i, tup) + inc i + + var stmtList = getAst(iterBody(iterCall)) + + var replaceCall = stmtList["ForStmt > StmtList > Call"] + + var yieldStmt = nnkYieldStmt.newTree( + nnkPar.newTree( + replaceCall[1], + replaceCall[2] + ) + ) + + stmtList["ForStmt > StmtList > Call"] = yieldStmt + iterDef[^1] = stmtList + + # Call the iterator surrounded by parens + result = nnkCall.newTree(nnkPar.newTree(iterDef)) + for arg in args: + result.add(arg) diff --git a/src/extensions/nim.cfg b/src/extensions/nim.cfg new file mode 100644 index 0000000..51df3ce --- /dev/null +++ b/src/extensions/nim.cfg @@ -0,0 +1 @@ +--path="../" diff --git a/src/extensions/nimnode.nim b/src/extensions/nimnode.nim new file mode 100644 index 0000000..1161ec6 --- /dev/null +++ b/src/extensions/nimnode.nim @@ -0,0 +1,85 @@ +import macros, strutils + +proc `[]=`*(node: NimNode, nodePath: string, value: NimNode) = + ## Set a NimNode based on a string search path like "StmtList > Command > Sym" + let nodeNames = nodePath.replace(" ", "").split(">") + + let targetIndex = len(nodeNames) - 1 + + var + index = -1 + targetParent: NimNode + stack: seq[seq[NimNode]] = @[@[node]] + indexList: seq[int] = newSeq[int](len(nodeNames)) + + + while stack.len() > 0 and index < targetIndex: + let parentList = stack.pop() + var newParents: seq[NimNode] = @[] + + block WhileBlock: + for parent in parentList: + inc index + for chIndex, child in parent.pairs: + let kindStr = ($child.kind)[3..^1] + + var tryParse = -1 + try: + tryParse = nodeNames[index].parseInt + except ValueError: + discard + + if kindStr == nodeNames[index] or tryParse == chIndex: + targetParent = parent + newParents.add(child) + indexList[index] = chIndex + + if index == targetIndex: + break WhileBlock + + if newParents.len() > 0: + stack.add(newParents) + + if index == targetIndex: + let modIndex = indexList.pop() + var curNode = node + for i in indexList: + curNode = curNode[i] + + curNode[modIndex] = value + + +proc `[]`*(node: NimNode, nodePath: string): NimNode = + ## Get a NimNode based on a string search path like "StmtList > Command > Sym" + let nodeNames = nodePath.replace(" ", "").split(">") + + let targetIndex = len(nodeNames) - 1 + + var + index = -1 + stack: seq[seq[NimNode]] = @[@[node]] + + while stack.len() > 0 and index < targetIndex: + let parentList = stack.pop() + var newParents: seq[NimNode] = @[] + + block WhileBlock: + for parent in parentList: + inc index + for i, child in parent.pairs: + let kindStr = ($child.kind)[3..^1] + + var tryParse = -1 + try: + tryParse = nodeNames[index].parseInt + except ValueError: + discard + + if kindStr == nodeNames[index] or tryParse == i: + result = child + newParents.add(child) + if index == targetIndex: + break WhileBlock + + if newParents.len() > 0: + stack.add(newParents)