WIP, translating Nim code to GLSL

shaders.nim

import glm
import macros
import sets
import strformat
import strutils
import tables

macro shader*(n: typed) =
  n.expectKind(nnkStmtList)
  let astFile = n.lineInfoObj.filename & ".ast"
  writeFile(astFile, n.treeRepr)

  proc recList(objectTy:NimNode):NimNode =
    objectTy.expectKind(nnkObjectTy)
    let recList = objectTy[2]
    recList.expectKind(nnkRecList)
    recList

  iterator fields(objectTy:NimNode):NimNode =
    let recList = objectTy.recList
    for rec in recList:
      case rec.kind
        of nnkIdentDefs:
          for i in 0..<rec.len-2:
            yield rec[i]
        of nnkSym:
          yield rec
        else: discard

  iterator vars(node:NimNode):NimNode =
    case node.kind:
      of nnkVarSection:
        let varSection = node
        for identDefs in varSection:
          identDefs.expectKind(nnkIdentDefs)
          for node in identDefs[0..<identDefs.len-2]:
            yield node
      of nnkIdentDefs:
        let identDefs = node
        for node in identDefs[0..<identDefs.len-2]:
          yield node
      else: error(&"cannot find variables in {node.kind}", node)

  template typeInst(node:NimNode):NimNode = node.getTypeInst

  proc resultTypeInst(procDef:NimNode):NimNode =
    for param in procDef.params:
      case param.kind:
        of nnkEmpty:
          return nil
        of nnkSym:
          return param.typeInst
        else: param.expectKind({nnkEmpty,nnkSym})

  iterator paramSyms(procDef:NimNode):NimNode =
    for param in procDef.params:
      case param.kind:
        of nnkEmpty:
          continue
        of nnkSym:
          continue
        of nnkIdentDefs:
          for node in param[0..<param.len-2]:
            case node.kind:
              of nnkEmpty:
                continue
              of nnkSym:
                yield node
              else: node.expectKind({nnkEmpty,nnkSym})
        else: param.expectKind({nnkEmpty,nnkSym})

  proc add[T](s:var HashSet[T], key:T):bool =
    if s.contains(key):
      result = false
    else:
      s.incl(key)
      result = true

  proc addUnique[T](s:var seq[T], item:T):bool =
    if s.contains(item):
      result = false
    else:
      s.add(item)
      result = true

  let glslNames = {
    "int"     :"long",
    "int32"   :"int",
    "int64"   :"long",
    "uint"    :"ulong",
    "uint32"  :"uint",
    "uint64"  :"ulong",
    "float32" :"float",
    "float64" :"double",
    "Vec2f"   :"vec2",
    "Vec3f"   :"vec3",
    "Vec4f"   :"vec4",
    "Mat4f"   :"mat4",
  }.toTable()

  proc getGlslName(typeName:string):string =
    result = typeName
    if glslNames.contains(typeName):
      result = glslNames[typeName]

  const glslGlobals = [
    "int",
    "long",
    "uint",
    "ulong",
    "float",
    "double",
    "vec2",
    "vec3",
    "vec4",
    "mat4",
  ]

  var globals:HashSet[string]

  var globalNodes:seq[NimNode]

  proc addTypeInst(glsl:var string, typeInst:NimNode) =
    if globalNodes.addUnique(typeInst):
      echo typeInst.treeRepr & " added"
    else:
      echo typeInst.treeRepr & " found"
    let typeName = getGlslName(typeInst.repr)
    if globals.add(typeName):
      let ty = typeInst.getType
      case ty.kind
        of nnkObjectTy:
          for field in ty.fields:
            let fieldTypeInst = field.typeInst
            glsl.addTypeInst(fieldTypeInst)
          glsl &= &"\nstruct {typeName} " & "{"
          for field in ty.fields:
            let fieldTypeInst = field.typeInst
            let fieldTypeName = getGlslName(fieldTypeInst.repr)
            glsl &= &"\n  {fieldTypeName} {field.repr};"
          glsl &= "\n};\n"
        else:
          glsl &= "\n/* " & typeName & "\n"
          glsl &= ty.treeRepr
          glsl &= "\n*/\n"

  proc addIndent(glsl:var string, depth:int) =
    for i in 0..<depth:
      glsl &= "  "

  proc addProcBody(glsl:var string, body:NimNode, depth:int=0) =
    case body.kind
      of nnkStmtList:
        let stmtList = body
        for node in stmtList:
          glsl.addIndent(depth)
          glsl.addProcBody(node)
          glsl &= ";\n"
      of nnkDiscardStmt:
        discard
      of nnkAsgn:
        let asgn = body
        let lhs = asgn[0]
        let rhs = asgn[1]
        glsl.addProcBody(lhs)
        glsl &= " = "
        glsl.addProcBody(rhs)
      of nnkDotExpr:
        let dotExpr = body
        let lhs = dotExpr[0]
        let rhs = dotExpr[1]
        glsl.addProcBody(lhs)
        glsl &= "."
        glsl.addProcBody(rhs)
      of nnkSym:
        let sym = body
        glsl &= sym.repr
      of nnkIntLit..nnkInt32Lit:
        let intLit = body
        glsl &= $intLit.intVal
      of nnkFloat32Lit:
        let floatLit = body
        glsl &= $floatLit.floatVal
      of nnkFloatLit,nnkFloat64Lit:
        let doubleLit = body
        glsl &= $doubleLit.floatVal & "lf"
      of nnkInfix:
        let infix = body
        let op = infix[0]
        let lhs = infix[1]
        let rhs = infix[2]
        glsl.addProcBody(lhs)
        glsl &= " " & op.repr & " "
        glsl.addProcBody(rhs)
      of nnkCall:
        let call = body
        let procName = call[0].repr
        glsl &= procName & "("
        for i in 1..<call.len:
          let arg = call[i]
          if not glsl.endsWith("("):
            glsl &= ", "
          glsl.addProcBody(arg)
        glsl &= ")"
      else:
        glsl &= "\n/* \n"
        glsl &= body.treeRepr
        glsl &= "\n*/\n"

  proc addProcDef(glsl:var string, procDef:NimNode) =
    let procName = procDef[0].repr
    if not globals.add(procName):
      return
    var resultTypeName = "void"
    var resultTypeInst = procDef.resultTypeInst
    if resultTypeInst != nil:
      resultTypeName = resultTypeInst.repr.getGlslName
      glsl.addTypeInst(resultTypeInst)
    for paramSym in procDef.paramSyms:
      let paramTypeInst = paramSym.typeInst
      glsl.addTypeInst(paramTypeInst)

    glsl &= &"\n{resultTypeName} {procName}("
    for paramSym in procDef.paramSyms:
      let paramTypeInst = paramSym.typeInst
      let paramTypeName = paramTypeInst.repr.getGlslName
      let paramName = paramSym.repr
      if not glsl.endsWith("("):
        glsl &= ", "
      glsl &= &"{paramTypeName} {paramName}"
    glsl &= ") {"

    if resultTypeInst != nil:
      glsl &= "\n"
      glsl.addIndent(depth=1)
      glsl &= &"{resultTypeName} result;"

    var body = procDef.body
    case body.kind:
      of nnkDiscardStmt:
        discard
      of nnkStmtList:
        glsl &= "\n"
        glsl.addProcBody(body,depth=1)
      else:
        glsl &= "\n"
        glsl.addProcBody(newStmtList(body),depth=1)

    if resultTypeInst != nil:
      glsl.addIndent(depth=1)
      glsl &= "return result;\n"
    glsl &= "}\n"

  proc addMain(glsl:var string, procDef:NimNode) =
    glsl.addProcDef(procDef)
    var resultTypeInst:NimNode
    let params = procDef.params
    for param in params:
      case param.kind
        of nnkEmpty:
          discard
        of nnkSym:
          let typeInst = param.typeInst
          glsl.addTypeInst(typeInst)
          resultTypeInst = typeInst
        of nnkIdentDefs:
          let sym = param[0]
          sym.expectKind(nnkSym)
          let typeInst = sym.typeInst
          glsl.addTypeInst(typeInst)
        else: param.expectKind({nnkEmpty,nnkSym,nnkIdentDefs})
    glsl &= "\nvoid main() {...}\n"

  type Stage = enum
    Vert,
    Frag,

  proc toGlsl(glsl:var string, node:NimNode, stage:Stage) =
    case stage
      of Vert:
        if globals.contains("vert"):
          return
      of Frag:
        if globals.contains("frag"):
          return
    case node.kind
      of nnkStmtList:
        for child in node:
          glsl.toGlsl(child, stage)
      of nnkTypeSection:
        let typeSection = node
        for typeDef in typeSection:
          let typeInst = typeDef[0]
          let typeName = getGlslName(typeInst.repr)
          case stage
            of Vert:
              discard
            of Frag:
              if typeName == "Vert":
                return
          glsl.addTypeInst(typeInst)
      of nnkVarSection:
        let varSection = node
        for varSym in varSection.vars:
          let typeInst = varSym.getTypeInst
          glsl.addTypeInst(typeInst)
          glsl &= &"\nuniform {typeInst.repr} {varSym.repr};\n"
      of nnkProcDef:
        let procDef = node
        let procName = procDef[0].repr
        case stage
          of Vert:
            if procName == "vert":
              glsl.addMain(procDef)
              return
            elif procName == "frag":
              return
            else:
              glsl.addProcDef(procDef)
          of Frag:
            if procName == "frag":
              glsl.addMain(procDef)
              return
            elif procName == "vert":
              return
            else:
              glsl.addProcDef(procDef)
      of nnkIncludeStmt:
        discard
      else:
        glsl &= "/*\n"
        glsl &= node.treeRepr
        glsl &= "*/\n"

  echo "-- vert --"
  block:
    var glsl = "#version 330\n"
    globals = glslGlobals.toHashSet
    glsl.toGlsl(n, Vert)
    echo glsl
  # echo "-- frag --"
  # block:
  #   var glsl = "#version 330\n"
  #   globals = glslGlobals.toHashSet
  #   glsl.toGlsl(n, Frag)
  #   echo glsl
  echo "----"


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

# it is useful to define uniform types outside of the shader code
type Transforms = object
  mvp:Mat4f

shader:
  var transforms:Transforms

  type
    Vert = object
      position:Vec4f
      color:Vec4f
    Frag = object
      position:Vec4f
      color:Vec4f

  proc f(v:Vec4f):Vec4f = v

  proc g() =
    discard

  proc vert*(v,w:Vert, #[vertId:uint, instId:uint]#):Frag =
    result.position = transforms.mvp * v.position.xyzw
    result.color = f(v.color)

  type Target = object
    color:Vec4f

  include pbr

  proc frag*(f:Frag):Target =
    result.color = f.color * pbr()