preserves-nim/src/preserves/private/preserves_schema_nim.nim

# SPDX-FileCopyrightText: 2021 ☭ Emery Hemingway
# SPDX-License-Identifier: Unlicense

import std/[os, strutils, tables]

import compiler/[ast, idents, renderer, lineinfos]

import ../../preserves, ../schemas

type TypeTable = OrderedTable[string, PNode]

proc add(parent, child: PNode): PNode {.discardable.} =
  parent.sons.add child
  parent

proc add(parent: PNode; children: varargs[PNode]): PNode {.discardable.} =
  parent.sons.add children
  parent

proc child(sn: SchemaNode): SchemaNode =
  assert(sn.nodes.len == 1)
  sn.nodes[0]

proc nn(kind: TNodeKind; children: varargs[PNode]): PNode =
  result = newNode(kind)
  result.sons.add(children)

proc nn(kind: TNodeKind; child: PNode): PNode =
  result = newNode(kind)
  result.sons.add(child)

proc ident(s: string): PNode =
  newIdentNode(PIdent(s: s), TLineInfo())

proc accQuote(n: PNode): Pnode =
  nkAccQuoted.newNode.add(n)

proc ident(sn: SchemaNode): PNode =
  var s: string
  case sn.kind
  of snkAlt:
    s = sn.altLabel.nimIdentNormalize
    s[0] = s[0].toLowerAscii
  of snkLiteral: s = $sn.value
  of snkRecord:
    s = sn.nodes[0].value.symbol
  of snkNamed:
    s = sn.name
  of snkDictionary, snkVariableTuple:
    s = "data"
  else:
    raiseAssert("no ident for " & $sn.kind & " " & $sn)
  s.ident.accQuote

proc typeIdent(sn: SchemaNode): PNode =
  case sn.kind
  of snkAtom:
    case sn.atom
    of akBool: ident"bool"
    of akFloat: ident"float32"
    of akDouble: ident"float64"
    of akInt: ident"BiggestInt"
    of akString: ident"string"
    of akBytes: nn(nkBracketExpr, ident"seq", ident"byte")
    of akSymbol: ident"string" # TODO: distinct string type for symbols?
  of snkNamed:
    sn.pattern.typeIdent
  of snkRef:
    var id = ident sn.refPath[sn.refPath.high]
    for i in countDown(sn.refPath.high.pred, 0):
      id = nn(nkDotExpr, ident(sn.refPath[i]), id)
    id
  else:
    ident"Preserve"

proc toExport(n: sink PNode): PNode =
  nkPostFix.newNode.add(ident"*", n)

proc newEmpty(): PNode = newNode(nkEmpty)

proc isConst(scm: Schema; sn: SchemaNode): bool =
  case sn.kind
  of snkLiteral: result = true
  of snkRef:
    if sn.refPath.len == 1:
      result = isConst(scm, scm.definitions[sn.refPath[0]])
  else: discard

proc literal(scm: Schema; sn: SchemaNode): Preserve =
  case sn.kind
  of snkLiteral: result = sn.value
  of snkRef:
    if sn.refPath.len == 1:
      result = literal(scm, scm.definitions[sn.refPath[0]])
    else:
      raiseAssert("not convertable to a literal: " & $sn)
  else:
    raiseAssert("not convertable to a literal: " & $sn)

proc isSymbolEnum(sn: SchemaNode): bool =
  if sn.kind == snkOr:
    for bn in sn.nodes:
      if bn.altBranch.kind != snkLiteral or
          bn.altBranch.value.kind != pkSymbol:
        return false
    result = true

proc toEnumTy(sn: SchemaNode): PNode =
  result = nkEnumTy.newNode.add newEmpty()
  for bn in sn.nodes:
    result.add bn.altLabel.nimIdentNormalize.ident.accQuote

proc toEnumDef(name: string; sn: SchemaNode): PNode =
  nkTypeDef.newNode.add(
    nkPragmaExpr.newNode.add(
      name.ident.toExport,
      nkPragma.newNode.add(ident"pure")),
      newEmpty(),
      sn.toEnumTy)

proc nimTypeOf(scm: Schema; known: var TypeTable; sn: SchemaNode; name = ""): PNode =
  case sn.kind
  of snkOr:
    if sn.isSymbolEnum:
      result = sn.toEnumTy
    else:
      let
        enumName = name.nimIdentNormalize & "Kind"
        enumIdent = ident(enumName)
      if enumName notin known:
        known[enumName] = toEnumDef(enumName, sn)
      let recCase = nkRecCase.newNode.add(
        nkIdentDefs.newNode.add(
          "kind".ident.toExport,
          enumName.ident,
          newEmpty()))
      for bn in sn.nodes:
        assert(bn.kind == snkAlt, $bn.kind)
        var recList = nkRecList.newNode
        case bn.altBranch.kind
        of snkRecord:
          case bn.altBranch.nodes.len
          of 0, 1: discard
          of 2:
            if not isConst(scm, bn.altBranch.nodes[1]):
              let label = bn.ident
              recList.add nkIdentDefs.newNode.add(
                label.toExport,
                nimTypeOf(scm, known, bn.altBranch.nodes[1], $label),
                newEmpty())
          else:
            for i, field in bn.altBranch.nodes:
              if i > 0 and (not isConst(scm, field)):
                let label = field.ident
                recList.add nkIdentDefs.newNode.add(
                  label.toExport,
                  nimTypeOf(scm, known, field, $label),
                  newEmpty())
        else:
          if isConst(scm, bn.altBranch):
            recList.add nkDiscardStmt.newNode.add(newEmpty())
          else:
            let label = bn.ident
            recList.add(nkIdentDefs.newNode.add(
              label,
              nimTypeOf(scm, known, bn.altBranch, $label),
              newEmpty()))
        let disc = nkDotExpr.newNode.add(
          enumIdent, bn.altLabel.nimIdentNormalize.ident.accQuote)
        if recList.len == 0:
          recList.add nn(nkDiscardStmt, newEmpty())
        recCase.add nkOfBranch.newNode.add(disc, recList)
      result = nn(nkRefTy, nn(nkObjectTy,
        newEmpty(),
        newEmpty(),
        nn(nkRecList, recCase)))
  of snkAny:
    result = ident"Preserve"
  of snkAtom:
    result = typeIdent(sn)
  of snkEmbedded:
    result = nimTypeOf(scm, known, sn.embed)
  of snkLiteral:
    result = case sn.value.kind # nearly verbatim from ../../preserves/src/preserves.nim
    of pkBoolean: ident"bool"
    of pkFloat: ident"float32"
    of pkDouble: ident"float64"
    of pkSignedInteger: ident"BiggestInt"
    of pkBigInteger: ident"BigInt"
    of pkString: ident"string"
    of pkByteString: nn(
      nkBracketExpr, ident"seq", ident"byte")
    of pkSymbol: ident"string"
    of pkRecord: ident"Preserve"
    of pkSequence: nn(
      nkBracketExpr, ident"seq", ident"Preserve")
    of pkSet: nn(
      nkBracketExpr, ident"HashSet", ident"Preserve")
    of pkDictionary: nn(
      nkBracketExpr, ident"Table", ident"Preserve", ident"Preserve")
  of snkSequenceOf:
    result = nkBracketExpr.newNode.add(
        ident"seq",
        nimTypeOf(scm, known, sn.child))
  of snkSetOf:
    result = nkBracketExpr.newNode.add(
        ident"HashedSet",
        nimTypeOf(scm, known, sn.child))
  of snkDictOf:
    result = nkBracketExpr.newNode.add(
        ident"Table",
        nimTypeOf(scm, known, sn.nodes[0]),
        nimTypeOf(scm, known, sn.nodes[1]))
  of snkRecord:
    case sn.nodes.len
    of 0, 1:
      result = nn(nkObjectTy,
        newEmpty(),
        newEmpty(),
        nn(nkRecList, nn(nkDiscardStmt, newEmpty())))
    else:
      let recList = nkRecList.newNode()
      for i, field in sn.nodes:
        if i > 0:
          let id = field.ident
          recList.add nkIdentDefs.newNode.add(
            id.toExport,
            nimTypeOf(scm, known, field, $id),
            newEmpty())
      result = nn(nkRefTy, nn(nkObjectTy,
        newEmpty(),
        newEmpty(),
        recList))
  of snkTuple:
    # TODO: the variable part
    result = nkTupleTy.newNode
    for tn in sn.nodes:
      if not isConst(scm, sn):
        result.add nkIdentDefs.newNode.add(
          tn.ident, nimTypeOf(scm, known, tn), newEmpty())
  of snkVariableTuple:
    result = nkTupleTy.newNode
    for i, tn in sn.nodes:
      if not isConst(scm, sn):
        if i == sn.nodes.high:
          result.add nkIdentDefs.newNode.add(
            tn.ident,
            nn(nkBracketExpr, ident"seq", nimTypeOf(scm, known, tn)),
            newEmpty())
        else:
          result.add nkIdentDefs.newNode.add(
            tn.ident, nimTypeOf(scm, known, tn), newEmpty())
  of snkDictionary:
    result = nkTupleTy.newNode
    for i in countup(0, sn.nodes.high, 2):
      let id = ident(sn.nodes[i+0])
      result.add nkIdentDefs.newNode.add(
        id,
        nimTypeOf(scm, known, sn.nodes[i+1], $id),
        newEmpty())
  of snkNamed:
    result = nimTypeOf(scm, known, sn.pattern, name)
  of snkRef:
    if sn.refPath.len == 1:
      let
        refName = sn.refPath[0]
        refDef = scm.definitions[refName]
      case refDef.kind
      of snkDictOf:
        result = nimTypeOf(scm, known, refDef, refName)
      else: result = typeIdent(sn)
    else:
      result = typeIdent(sn)
  else:
    result = nkCommentStmt.newNode
    result.comment.add("Missing type generator for " & $sn.kind & " " & $sn)

proc exportIdent(id: string): PNode = nn(nkPostFix, ident"*", ident(id))

proc generateConstProcs(result: var seq[PNode]; name: string; def: SchemaNode) =
  case def.kind
  of snkLiteral:
    var stmts = nn(nkStmtList)
    case def.value.kind
    of pkSignedInteger:
      discard stmts.add newIntNode(nkIntLit, def.value.int)
    of pkSymbol:
      discard stmts.add nn(nkCall,
        ident"symbol",
        PNode(kind: nkStrLit, strVal: def.value.symbol))
    else:
      raiseAssert("conversion of " & $def & " to a Nim literal is not implemented")
    var procId = name
    procId[0] = procId[0].toLowerAscii
    let constProc= nn(nkProcDef,
      exportIdent(procId),
      newEmpty(),
      newEmpty(),
      nn(nkFormalParams, ident"Preserve"),
      newEmpty(),
      newEmpty(),
      stmts)
    constProc.comment = "``" & $def & "``"
    result.add constProc
  else: discard

proc toNimLit(sn: SchemaNode): PNode =
  assert(sn.kind == snkLiteral, $sn)
  case sn.value.kind
  of pkSymbol:
    nkCall.newNode.add(
      ident"symbol",
      PNode(kind: nkStrLit, strVal: sn.value.symbol))
  else:
    raiseAssert("no Nim literal for " & $sn)

proc literalToPreserveCall(pr: Preserve): PNode =
  var prConstr = nn(nkObjConstr, ident"Preserve")
  proc constr(kind, field: string; lit: PNode) =
    prConstr.add nn(nkExprColonExpr, ident"kind", ident(kind))
    prConstr.add nn(nkExprColonExpr, ident(field), lit)
  case pr.kind
  of pkBoolean:
    constr($pr.kind, "bool", if pr.bool: ident"true" else: ident"false")
  of pkFloat:
    constr($pr.kind, "float", newFloatNode(nkFloat32Lit, pr.float))
  of pkDouble:
    constr($pr.kind, "double", newFloatNode(nkFloat64Lit, pr.double))
  of pkSignedInteger:
    constr($pr.kind, "int", newIntNode(nkInt64Lit, pr.int))
  of pkString:
    constr($pr.kind, "string", newStrNode(nkTripleStrLit, pr.string))
  of pkByteString:
    return nn(nkCall, ident"parsePreserves", newStrNode(nkTripleStrLit, $pr))
  of pkSymbol:
    constr($pr.kind, "symbol", newStrNode(nkStrLit, pr.symbol))
  else:
    raise newException(ValueError, "refusing to convert to a literal: " & $pr)
  prConstr

proc tupleConstructor(scm: Schema; sn: SchemaNode; ident: PNode): Pnode =
  let seqBracket = nn(nkBracket)
  for i, field in sn.nodes:
    if isConst(scm, field):
      seqBracket.add literalToPreserveCall(literal(scm, field))
    elif sn.kind == snkTuple or i < sn.nodes.high:
      seqBracket.add nn(nkCall,
        ident"toPreserve",
        nn(nkDotExpr, ident, field.ident))
  let seqConstr = nn(nkPrefix, ident"@", seqBracket)
  let colonExpr = nn(nkExprColonExpr, ident"sequence")
  if sn.kind == snkTuple:
    colonExpr.add seqConstr
  else:
    colonExpr.add nn(nkInfix,
      ident"&",
      seqConstr,
      nn(nkDotExpr,
        nn(nkCall, ident"toPreserve",
          nn(nkDotExpr,
            ident, sn.nodes[sn.nodes.high].ident)),
        ident"sequence"))
  nn(nkObjConstr,
    ident"Preserve",
    nn(nkExprColonExpr, ident"kind", ident"pkSequence"),
    colonExpr)

proc generateProcs(result: var seq[PNode]; scm: Schema; name: string; sn: SchemaNode) =
  case sn.kind
  of snkOr:
    var
      enumId = name.ident
      paramId = name.toLowerAscii.ident.accQuote
      orStmts = nn(nkStmtList)
    if sn.isSymbolEnum:
      let caseStmt = nn(nkCaseStmt, paramId)
      for bn in sn.nodes:
        caseStmt.add nn(nkOfBranch,
          nn(nkDotExpr,
            enumId,
            bn.altLabel.nimIdentNormalize.ident.accQuote),
          nn(nkCall,
            ident"symbol", PNode(kind: nkStrLit, strVal: $bn.altLabel)))
      orStmts.add caseStmt
    else:
      let caseStmt = nn(nkCaseStmt, nn(nkDotExpr, paramId, ident"kind"))
      proc genStmts(stmts: PNode; fieldId: PNode; sn: SchemaNode) =
        case sn.kind
        of snkLiteral:
           stmts.add literalToPreserveCall(literal(scm, sn))
        of snkOr, snkRecord, snkRef:
          if sn.kind == snkRef and sn.refPath.len == 1:
            let refDef = scm.definitions[sn.refPath[0]]
            genStmts(stmts, fieldId, refDef)
          else:
            stmts.add nn(nkCall,
              ident"toPreserve",
              nn(nkDotExpr, paramId, fieldId))
        of snkTuple, snkVariableTuple:
          stmts.add tupleConstructor(scm, sn, nn(nkDotExpr, paramId, fieldId))
        else:
          raiseAssert("no case statement for " & $sn.kind & " " & $sn)
      for bn in sn.nodes:
        let stmts = nn(nkStmtList)
        genStmts(stmts, bn.ident, bn.altBranch)
        caseStmt.add nn(nkOfBranch,
          nn(nkDotExpr,
            ident(name & "Kind"),
            bn.altLabel.nimIdentNormalize.ident.accQuote),
          stmts)
      orStmts.add caseStmt
    result.add nn(nkProcDef,
      exportIdent("toPreserveHook"),
      newEmpty(),
      newEmpty(),
      nn(nkFormalParams,
        ident"Preserve",
        nn(nkIdentDefs,
          paramId, ident(name), newEmpty())),
      newEmpty(),
      newEmpty(),
      orStmts)
  of snkRecord:
    var
      params = nn(nkFormalParams, ident"Preserve")
      initRecordCall = nn(nkCall,
        ident"initRecord",
        sn.nodes[0].toNimLit)
    for i, field in sn.nodes:
      if i > 0:
        let
          id = field.ident
        var fieldType  = field.typeIdent
        if fieldType.kind != nkIdent or fieldType.ident.s != "Preserve":
          fieldType =
            nn(nkInfix,
              ident"|",
              fieldType,
              ident"Preserve")
        params.add nn(nkIdentDefs,
          id, fieldType, newEmpty())
        initRecordCall.add(
          nn(nkCall, ident"toPreserve", id))
    var procId = name
    procId[0] = procId[0].toLowerAscii
    result.add nn(nkProcDef,
      exportIdent(procId),
      newEmpty(),
      newEmpty(),
      params,
      newEmpty(),
      newEmpty(),
      nn(nkStmtList,
        PNode(
          kind: nkCommentStmt,
          comment: "Preserves constructor for ``" & name & "``."),
        initRecordCall))
    block:
      let paramId = name.toLowerAscii.ident.accQuote
      initRecordCall = nn(nkCall,
        ident"initRecord",
        sn.nodes[0].toNimLit)
      for i, field in sn.nodes:
        if i > 0:
          initRecordCall.add nn(nkCall,
            ident"toPreserve",
            nn(nkDotExpr, paramId, field.ident))
      result.add nn(nkProcDef,
        exportIdent("toPreserveHook"),
        newEmpty(),
        newEmpty(),
        nn(nkFormalParams,
          ident"Preserve",
          nn(nkIdentDefs,
            paramId, ident(name), newEmpty())),
        newEmpty(),
        newEmpty(),
        nn(nkStmtList, initRecordCall))
  of snkTuple, snkVariableTuple:
    let paramId = name.toLowerAscii.ident.accQuote
    result.add nn(nkProcDef,
      exportIdent("toPreserveHook"),
      newEmpty(),
      newEmpty(),
      nn(nkFormalParams,
        ident"Preserve",
        nn(nkIdentDefs,
          paramId, ident(name), newEmpty())),
      newEmpty(),
      newEmpty(),
      nn(nkStmtList, tupleConstructor(scm, sn, paramId)))
  else: discard

proc collectRefImports(imports: PNode; sn: SchemaNode) =
  case sn.kind
  of snkLiteral:
    if sn.value.isDictionary:
      imports.add ident"std/tables"
  of snkDictOf:
    imports.add ident"std/tables"
  of snkRef:
    if sn.refPath.len > 1:
      imports.add ident(sn.refPath[0])
  else:
    for child in sn.items:
      collectRefImports(imports, child)

proc collectRefImports(imports: PNode; scm: Schema) =
  if scm.embeddedType.contains {'.'}:
    let m = split(scm.embeddedType, '.', 1 )[0]
    imports.add ident(m)
  for _, def in scm.definitions:
    collectRefImports(imports, def)

proc generateNimFile*(scm: Schema; path: string) =
  var
    knownTypes: TypeTable
    typeSection = newNode nkTypeSection
    procs: seq[PNode]
    megaType: PNode
  for name, def in scm.definitions.pairs:
    if isConst(scm, def):
      generateConstProcs(procs, name, def)
    else:
      var name = name
      name[0] = name[0].toUpperAscii
      if megaType.isNil:
        megaType = ident(name)
      else:
        megaType = nn(nkInfix,
          ident"|", megaType, ident(name))
      let t = nimTypeOf(scm, knownTypes, def, name)
      t.comment = "``" & $def & "``"
      case def.kind
      of snkAtom:
        knownTypes[name] = nkTypeDef.newNode.add(
          name.ident.toExport, newEmpty(), t)
      else:
        if def.kind == snkRecord:
          knownTypes[name] = nn(nkTypeDef,
            nn(nkPragmaExpr,
              name.ident.toExport,
              nn(nkPragma,
                nn(nkExprColonExpr,
                  ident"record",
                  PNode(kind: nkStrLit, strVal: $def.nodes[0].value.symbol)))),
              newEmpty(),
              t)
        else:
          case t.kind
          of nkEnumTy:
            knownTypes[name] = nn(nkTypeDef,
              nn(nkPragmaExpr,
                name.ident.toExport,
                nn(nkPragma, ident"pure")),
                newEmpty(),
                t)
          of nkDistinctTy:
            knownTypes[name] = nn(nkTypeDef,
              nn(nkPragmaExpr,
                name.ident.toExport,
                nn(nkPragma,
                  nn(nkExprColonExpr,
                    ident"borrow",
                    accQuote(ident".")))),
                newEmpty(),
                t)
          else:
            knownTypes[name] = nn(nkTypeDef,
              name.ident.toExport, newEmpty(), t)
      generateProcs(procs, scm, name, def)
  for typeDef in knownTypes.values:
    typeSection.add typeDef
  var imports = nkImportStmt.newNode.add(
    ident"std/typetraits",
    ident"preserves")
  collectRefImports(imports, scm)
  procs.add nn(nkProcDef,
      "$".ident.accQuote.toExport,
      newEmpty(),
      newEmpty(),
      nn(nkFormalParams,
        ident"string",
        nn(nkIdentDefs,
          ident"x",
          megaType,
          newEmpty())),
      newEmpty(),
      newEmpty(),
      nn(nkStmtList,
        nn(nkCall, ident"$", nn(nkCall, ident"toPreserve", ident"x"))))
  procs.add nn(nkProcDef,
      "encode".ident.accQuote.toExport,
      newEmpty(),
      newEmpty(),
      nn(nkFormalParams,
        nn(nkBracketExpr, ident"seq", ident"byte"),
        nn(nkIdentDefs,
          ident"x",
          megaType,
          newEmpty())),
      newEmpty(),
      newEmpty(),
      nn(nkStmtList,
        nn(nkCall, ident"encode", nn(nkCall, ident"toPreserve", ident"x"))))
  var module = newNode(nkStmtList).add(
      imports,
      typeSection
    ).add(procs)
  writeFile(path, renderTree(module, {renderNone, renderIr}))

when isMainModule:
  import std/parseopt
  var inputs: seq[string]
  for kind, key, val in getopt():
    case kind
    of cmdLongOption:
      case key
      else: quit("unhandled option " & key)
    of cmdShortOption:
      case key
      else: quit("unhandled option " & key)
    of cmdArgument:
      inputs.add absolutePath(key)
    of cmdEnd: discard
  for filepath in inputs:
    let
      scm = parsePreservesSchema(readFile filepath, filepath)
      (dir, name, _) = splitFile(filepath)
      outputPath = dir / name & ".nim"
    generateNimFile(scm, outputPath)
    stdout.writeLine(outputPath)