either.nim

## Introduces a number of `Either` types which are union-like types which can hold one single
## value of multiple possible types. They are meant to be used as parts of types autogenerated
## by schema parsers, for example JSON where a valid property may be a `string` or `int`.
## 
## Generates types `Either2[T0, T1]` to `Either[T0, T1, T2, T3, T4, T5, T6, T7]`.

import
  macros,
  options,
  sets,
  strutils,
  typetraits

# Helper functions for producing a easily human-readable string
# from a type node, for example `Either3[SomeType, string, bool]`.
macro `$!`(t: untyped): string = t.getTypeInst().repr
proc `$!`(t: NimNode): string = t.getTypeInst().repr

macro makeEitherType(size: static[int]): untyped =
  let size = size # https://github.com/nim-lang/Nim/issues/9509
  
  proc process(node: NimNode, index = -1): NimNode =
    case node.kind:
      of nnkIdent:
        return case $node.ident:
          # Replace "X" in `X`, `valueX` and `TX` with the index
          # of the iteration inside `case`, `when` or `set` proc.
          of      "X": newLit(index)
          of "valueX": newIdentNode("value" & $index)
          of     "TX": newIdentNode("T" & $index)
          # Replace "SIZE" in `EitherSIZE` and `SIZE_MINUS_ONE` with the
          # size of the created Either type (in the latter case, one less).
          of     "EitherSIZE": newIdentNode("Either" & $size)
          of "SIZE_MINUS_ONE": newLit(size - 1)
          else: node
      
      # Replace `[TYPES]` in type definition and usage with `[T0, T1, ...]`.
      # That is, `type EitherSIZE*[TYPES]`, `EitherSIZE[TYPES]` and `@[TYPES]`.
      of nnkIdentDefs, nnkBracketExpr, nnkBracket:
        # In BracketExpr, the first node is `EitherSIZE`, so skip that.
        let first = if node.kind == nnkBracketExpr: 1 else: 0
        if (node[first].eqIdent("TYPES")):
          node.del(first)
          for i in 0..<size:
            node.insert(first + i, newIdentNode("T" & $i))
      
      # Expand case and when statements with nodes duplicated times `size`.
      # These duplicated nodes will have their `index` assigned.
      of nnkRecCase, nnkCaseStmt, nnkWhenStmt:
        # In RecCase and CaseStmt, the first node is is the `either.which` part, so skip that.
        let first = if node.kind != nnkWhenStmt: 1 else: 0
        for i in 1..<size:
          node.insert(first + i, process(node[first].copyNimTree(), i))
        node[first] = process(node[first], 0)
      
      else: discard
    
    # Recursively process all the children of this node.
    for i in 0..<node.len:
      node[i] = process(node[i], index)
    
    node
  
  # Instead of creating the entire AST from scratch using `newTree` and such, we parse a
  # string and replace, modify or duplicate specific nodes to do what we want. This makes
  # it a bit easier to see the general makeup and which procs are available.
  
  # Not using `quote` because it changes the AST in ways that make things more difficult.
  
  # Now just recursively "process" the AST, replacing and manipulating parts of it.
  # (TIP: Remove the triple-quotes to see syntax formatting and make this easier to read.)
  result = process parseStmt """
type
  EitherSIZE*[TYPES] = object
    case which: range[0..SIZE_MINUS_ONE]
      of X: valueX: TX

proc getTypeStr*[TYPES](either: EitherSIZE[TYPES]): string =
  case either.which:
    of X: $TX

template has*[TYPES](either: EitherSIZE[TYPES], want: typedesc): bool =
  when want is TX: either.which == X
  else: {.fatal: ("'$#' is not a valid type for $#" % [$want, $!either]).}

template unsafeGet*[TYPES](either: EitherSIZE[TYPES], want: typedesc): untyped =
  when want is TX: either.valueX
  else: {.fatal: ("'$#' is not a valid type for $#" % [$want, $!either]).}

template get*[TYPES](either: EitherSIZE[TYPES], want: typedesc): untyped =
  if either.has(want): either.unsafeGet(want)
  else: raise newException(ValueError, "$# is not a '$#'" % [$!either, $want])

template option*[TYPES](either: EitherSIZE[TYPES], want: typedesc): Option[want] =
  if either.has(want):
    when want is TX: either.valueX.some
    # No `else` case needed, `has` will already error on invalid type.
  else: want.none

template set*[TYPES](either: var EitherSIZE[TYPES], value: typed) =
  when type(value) is TX:
    either.which  = X
    either.valueX = value
  else: {.fatal: ("'$#' is not a valid type for $#" % [$want, $!either]).}
"""

makeEitherType(2)
makeEitherType(3)
makeEitherType(4)
makeEitherType(5)
makeEitherType(6)
makeEitherType(7)
makeEitherType(8)


template `[]`*(either: typed, want: typedesc): untyped =
  either.get(want)

template `[]=`*(either: typed, want: typedesc, value: want): untyped =
  either.set(value)


macro lookupIndex(either: typed, expectedType: typed): int =
  let recCase = either.getType()[2][0]
  for i in 1..<recCase.len:
    let ofType = recCase[i][1].getTypeInst()
    # FIXME: Does not account for aliases. `sameType` doesn't appear to work, either. Maybe it only does for `typedesc`.
    if ofType == expectedType: return i - 1
  # FIXME: Somehow pass the source node of `expectedType` into `error` to get errors to appear at the right place.
  error("'$#' is not a valid type for $#" % [expectedType.repr, $!either])

macro fillBranchContents(either: typed, valueIdent: untyped, branches: untyped, protoCaseStmt: typed): untyped =
  # Fills in branches from `prototypeCaseStmt` with `valueIdent` definition and actual content.
  # 
  #   case either.which:
  #     of 0: discard
  #     of 1: discard
  #     [...]
  # 
  # is transformed into:
  # 
  #   case either.which:
  #     of 0:
  #       let valueIdent = either.value0
  #       [branch content]
  #     of 1:
  #       let valueIdent = either.value1
  #       [branch content]
  #     [...]
  
  result = protoCaseStmt
  for i in 1..<protoCaseStmt.len:
    var protoBranch = protoCaseStmt[i]
    if protoBranch.kind != nnkOfBranch: continue
    let which = protoBranch[0].intVal # The looked-up index.
    
    var contents = branches[i - 1][1]
    protoBranch[1] = contents
    
    let valueX = newIdentNode("value" & $which)
    contents.insert(0): quote do:
      let `valueIdent` = `either`.`valueX`

macro match*(either: typed, valueIdent: untyped, branches: varargs[untyped]): untyped =
  let either = either # https://github.com/nim-lang/Nim/issues/9509
  valueIdent.expectKind(nnkIdent)
  
  # The `match` macro will create case statement like so:
  # 
  #   case either.which:
  #     of lookupIndex(either, type1): discard
  #     of lookupIndex(either, type2): discard
  #     [...]
  #     else: [branch content]
  # 
  # The branch contents (expect for the else branch) are deliberately omitted, because
  # they likely make use of `valueIdent`, which is not definited at that point.
  # However, `lookupIndex` needs to be resolved, so the unfinished case statement is
  # passed to `fillBranchContents` as typed, and the `branches` as untyped, where the
  # contents are actually filled in.
  
  let caseStmt = quote do:
    case `either`.which:
  
  for branch in branches:
    if branch.kind == nnkOfBranch:
      let expectedType = branch[0]
      # We will strip `case nil` after. An `of` branch without it is not valid in `quote`.
      # Simply use `discard` as the branch content for now. This way we can force the entire
      # case statement to be typed, resolving the lookup index. Content will be added later.
      let caseStmtSingle = quote do:
        case nil:
          of lookupIndex(`either`, `expectedType`): discard
      caseStmt.add(caseStmtSingle[1])
    else: caseStmt.add(branch)
  
  result = quote do:
    fillBranchContents(`either`, `valueIdent`):
      `branches`
    do:
      `caseStmt`


when isMainModule:
  var e2: Either2[string, Option[bool]]
  e2.set(true.some)
  
  echo "e2 is: " & $e2.getTypeStr()
  echo "e2 has string: " & $e2.has(string)
  echo "e2 has bool?: " & $e2.has(Option[bool])
  # Error: 'bool' is not a valid type for 'Either2[string, Option[bool]]'
  # echo "e2 has int?: " & $e2.has(bool)
  
  echo $e2.option(string)
  echo $e2.get(Option[bool])
  echo $e2[Option[bool]]
  # Error: 'bool' is not a valid type for 'Either2[string, Option[bool]]'
  # echo $e2.option(bool)
  
  e2.match(value):
  of string: echo "IS STRING => \"", value, "\""
  of Option[bool]: echo "IS BOOL? => ", $value
  # Error: 'bool' is not a valid type for 'Either2[string, Option[bool]]'
  # of bool: echo "IS INT! => ", $value
  else: echo "IS NOT STRING!" # Shouldn't be possible, all branches are satisfied.