Nim `impl` macro

gistfile1.nim

import macros

# ---------- ---------- ---------- #


proc hasPragma(procedure:expr, name:string): bool {.compileTime.} =
  # checks if a procedure has a specific pragma
  let procPragma = procedure[4]#.pragma
  for p in procPragma.children:
    if p.repr == name:
      return true

# ---

proc setupParams(procedure:expr, initParams, callParams:var seq[PNimrodNode]) {.compileTime.} =
  # adds params to an initilizer or constructor
  let procParams = procedure[3]#.params
  let thisName = procParams[1][0]
  let typeIdent = procParams[1][1]
  
  initParams.add(typeIdent)
  initParams.add(newIdentDefs(thisName, newNimNode(nnkTypeOfExpr).add(typeIdent)))
  callParams.add(newIdentNode("result"))
  
  for i in 2 .. <procParams.len:
    initParams.add(procParams[i])
    for c in 0 .. procParams[i].len-3:
      callParams.add(procParams[i][c])


# --- --- --- #


macro allocator*(procedure:expr{nkProcDef,nkTemplateDef,nkMacroDef}): stmt {.immediate.} =
  ## Converts a procedure into an allocator.
  ##
  ## example:
  ##   proc foo(this:Bar, ...) {.allocator, custom.}
  ##
  ## converts to this:
  ##   proc foo(this:type Bar, params...): Bar {.custom.} =
  ##     template construct(this:Bar, params...) =
  ##       body...
  ##     construct(result, params...)
  ##
  let noInit = procedure.hasPragma("noinit")
  let consName = newIdentNode("construct")
  var initParams = newSeq[PNimrodNode]()
  var callParams = newSeq[PNimrodNode]()
  
  procedure.setupParams(initParams, callParams)
  
  let constructor = newNimNode(nnkTemplateDef).add(
    consName,
    newEmptyNode(),
    newEmptyNode(),
    procedure[3],#.params
    newEmptyNode(),
    newEmptyNode(),
    procedure.body
  )
  
  let initBody = newStmtList().add(constructor)
  if not noInit:
    initBody.add newCall(newDotExpr(ident"system", ident"new"), ident"result")
  initBody.add newCall(consName, callParams)
  
  result = newProc(procedure.name, initParams, initBody)
  result.pragma = copy(procedure.pragma)
  if not noInit:
    result.pragma.add(ident"noinit")


# --- --- --- #


macro constructor*(procedure:expr{nkProcDef,nkTemplateDef,nkMacroDef}): stmt {.immediate.} =
  ## Converts a procedure into an constructor.
  ##
  ## example:
  ##   proc foo(this:Bar, params...) {.constructor, custom.} =
  ##     body...
  ##
  ## converts to this:
  ##   proc construct(this:Bar, params...) {.custom.} =
  ##     body...
  ##   proc foo(this:type Bar, params...): Bar {.inline noinit custom.} =
  ##     system.new(result)
  ##     construct(result, params...)
  ##
  let initName = procedure.name
  let consParams = procedure[3]#.params
  var initParams = newSeq[PNimrodNode]()
  var callParams = newSeq[PNimrodNode]()
  
  procedure.setupParams(initParams, callParams)
  
  let constructor = newNimNode(nnkProcDef).add(
    initName,
    newEmptyNode(),
    newEmptyNode(),
    consParams,
    procedure.pragma,
    newEmptyNode(),
    procedure.body
  )
  
  let initBody = newStmtList().add(
    newCall(newDotExpr(ident"system", ident"new"), ident"result"),
    newCall(initName.basename, callParams)
  )
  
  let allocator = newProc(initName, initParams, initBody, procedure.kind)
  allocator.pragma = copy(procedure.pragma)
  allocator.pragma.addIdentIfAbsent("inline")
  allocator.pragma.addIdentIfAbsent("noinit")
  
  return newStmtList().add(constructor, allocator)


# ---------- ---------- ---------- #


const
  # default name of first param
  thisName = "this"
  
  # set of node kinds which impliment a type
  thisProcKinds = {
    nnkProcDef,
    nnkTemplateDef,
    nnkMacroDef,
    nnkIteratorDef,
    nnkConverterDef
  }


# --- --- --- #


iterator pairs(node:PNimrodNode): tuple[i:int, n:PNimrodNode] =
  for i in 0 .. <node.len:
    yield (i, node[i])


# --- --- --- #


macro impl*(head:expr{nkIdent}, body:stmt): stmt {.immediate.} =
  ## Impliments a set of procedures for a type.
  ##
  ## example:
  ##   impl int:
  ##     proc foo = ...
  ## 
  ## converts into this:
  ##   proc foo(this:int) = ...
  ##
  for i, node in body:
    case node.kind:
      
      of thisProcKinds:
        let params = node[3]#.params
        params.insert(1, newIdentDefs(ident(thisName), head))
      
      of nnkCall:
        if node[0].repr == "global":
          body[i] = newNimNode(nnkCommand).add(ident"impl", newNimNode(nnkTypeOfExpr).add(head), node[1])
      
      of nnkCommand:
        discard # TODO: ...
      
      else:
        discard
  
  return body


# ---------- ---------- ---------- #


proc extractTypeName(node:PNimrodNode): PNimrodNode {.compileTime noinit.} =
  assert node.kind == nnkInfix
  case node[0].repr:
    of "as":
      return node[1]
    of "*":
      return postfix(node[1], "*")
    of "of":
      return extractTypeName(node[1])
    else:
      quit "Invalid infix: " & node[0].repr

# ---

proc extractTypeStruct(node:PNimrodNode): PNimrodNode {.compileTime noinit.} =
  assert node.kind in {nnkInfix, nnkPrefix}
  case node[0].repr:
    of "as":
      return node[2]
    of "*":
      assert node[2].kind == nnkPrefix and node[2][0].repr == "as"
      return node[2][1]
    of "of":
      return extractTypeStruct(node[1])
    else:
      quit "Invalid infix/prefix: " & node[0].repr

# ---

proc extractTypeParent(node:PNimrodNode): PNimrodNode {.compileTime noinit.} =
  assert node.kind == nnkInfix
  if node[0].repr == "of":
    return newNimNode(nnkOfInherit).add(node[2])
  else:
    return newEmptyNode()
    #return extractTypeStruct(node[2])

# ---

proc extractTypeObject(structNode:PNimrodNode): PNimrodNode {.compileTime noinit.} =
  if structNode.kind != nnkObjectTy:
    assert structNode.kind in {nnkRefTy, nnkPtrTy}
    return extractTypeObject(structNode[0])
  else:
    return structNode


# --- --- --- #


macro class*(head:expr{nkInfix,nkPrefix}, body:stmt): stmt {.immediate.} =
  ## Impliments a general purpose OOP class-like type and procedures.
  ##
  ## example:
  ##
  ##   class Person* as ref object:
  ##     var
  ##       name: string
  ##       age: int
  ##     
  ##     proc new*(name:string, age = 0) {.allocator.} =
  ##       this.name = name
  ##       this.age = age
  ##     
  ##     proc greet* =
  ##       ...
  ##
  ## converts into this:
  ##
  ##   type Person* = ref object:
  ##     name: string
  ##     age: int
  ##   
  ##   proc new*(this:type Person, name:string, age=0): Person =
  ##     template constructor(this:Person, name:string, age=0) =
  ##       this.name = name
  ##       this.age = age
  ##     system.new(result)
  ##     constructor(result, name, age)
  ##   
  ##   proc greet*(this:Person) =
  ##     ...
  ##
  let typeName   = extractTypeName(head)
  let typeStruct = extractTypeStruct(head)
  let typeParent = extractTypeParent(head)
  let typeObject = extractTypeObject(typeStruct)
  let typeRecs   = newNimNode(nnkRecList)
  let implNodes  = newStmtList()
  
  # handle body
  for node in body.children:
    case node.kind:
      
      of nnkVarSection:
        # add each var to type records
        for i in 0 .. <node.len:
          # TODO: handle default assignment condition
          typeRecs.add(node[i])
      
      of nnkCall, nnkCommand, thisProcKinds:
        # add node to be implimented
        implNodes.add(node)
      
      else:
        discard
  
  # add type records
  typeObject.add(
    newEmptyNode(),
    typeParent,
    typeRecs
  )
  
  # compose final statement
  return newStmtList().add(
    
    # add the type
    newNimNode(nnkTypeSection).add(
      newNimNode(nnkTypeDef).add(
        typeName,
        newEmptyNode(),
        typeStruct
      )
    ),
    
    # impl the body
    newCall(ident"impl", typeName.basename)
      .add(implNodes)
  )