ch35.md 4.0 KB
Chapter 35. Collect macro
As our first complex example we will look at how Nim’s collect macro can be implemented. The standard library already contains collect, it can be found in std/sugar. collect is the preferred method of turning a potentially nested loop construct from a statement to an expression. Instead of: import std / tables const Data = toTable({"a": 1, "b": 2, "c": 3}) var s = newSeq[string]() for k, v in Data.pairs:
if v mod 2 == 0:
s.add k
One can use the more declarative: import std / [tables, sugar] const Data = toTable({"a": 1, "b": 2, "c": 3}) let s = collect(newSeq):
for k, v in Data.pairs:
if v mod 2 == 0: k
An an exercise we will reimplement collect. For a beginner, writing a macro is usually a hard task. As the first step we postulate the code pattern that the macro needs to expand to: collect(constructorCall): body should be translated into something like:
235
block:
var tmp = constructorCall[typeOf(body)]()
sinkInto(body, tmp.add)
tmp
where sinkInto(body, tmp.add) describes the AST where the final expression x of body is replaced by tmp.add x. We have to walk if-expressions, loops and “statement list expressions” to arrive at the “final expression” which is the part of the body that produces the value: import macros proc sinkInto(n, fullBody, res, bracketExpr: NimNode): NimNode = 1
case n.kind
of nnkStmtList, nnkStmtListExpr, nnkBlockStmt, nnkBlockExpr,
nnkWhileStmt, nnkForStmt, nnkElifBranch, nnkElse, nnkElifExpr,
nnkOfBranch, nnkExceptBranch: 2
result = copyNimTree(n)
if n.len >= 1:
result[^1] = sinkInto(n[^1], fullBody, res, bracketExpr)
of nnkIfExpr, nnkIfStmt, nnkTryStmt: 3
result = copyNimTree(n)
for i in 0..<n.len:
result[i] = sinkInto(n[i], fullBody, res, bracketExpr)
of nnkCaseStmt:
result = copyNimTree(n)
for i in 1..<n.len: 4
result[i] = sinkInto(n[i], fullBody, res, bracketExpr)
else:
if bracketExpr.len == 1: 5
bracketExpr.add(newCall(bindSym"typeof", fullBody))
result = newCall(bindSym"add", res, n) 6
macro collect*(init, body: untyped): untyped =
let res = genSym(nskVar, "collectResult") 7
let bracketExpr = newTree(nnkBracketExpr, init) 8
let transformedBody = sinkInto(body, body, res, bracketExpr) 9
let call = newTree(nnkCall, bracketExpr) 10
result = newTree(nnkStmtListExpr, newVarStmt(res, call),
transformedBody, res) 11
1 Traverses n recursively and produces a copy of n except that the value
producing subexpression x is replaced by add(res, x). fullBody is the full
body and it is used for producing init[typeof(body)]() which is
accomplished by modifying bracketExpr.
236 2 For nnkStmtListExpr and the like we only follow the last child. The last child can be accessed via n[^1]. 3 For if expressions and the like we follow all possible branches. This allows for code like if cond: a else: b to be transformed into: if cond: add(tmp, a) else: add(tmp, b). 4 A case expression is just like an if expression except that we start from 1 here in order to skip the selection expression which is not the value producing expression that we are interested in. 5 If the bracketExpr is still the init expression, add typeof(body) to it producing init[typeof(body)]. 6 We arrived at the value producing expression. Transform it to add(res, value). 7 We create a (var res = init[typeof(body)]; transformedBody; res) construct which is called an nnkStmtListExpr tree. res is a fresh variable produced from macros.genSym. 8 The construct init[T] is generated as an nnkBracketExpr. 9 Let sinkInto perform the recursive traversal. 10 We transform init[T] to init[T](). 11 The result of the macro is this (var res = init[typeof(body)]; transformedBody; res) construct. Note that this is a simplified implementation, the standard library’s collect macro implements more features and is more flexible.