callOperator_experiment.nim

type
  Function = ref object
    stages: seq[Stage]
    output: Parameter
    schedule: FunctionSchedule

  Call = ref object
    function: Function
    params: seq[LuxNode]

  Stage = ref object
    params: seq[LuxNode]
    definition: LuxNode

    condition: LuxNode
      ## wrapped in a
      ## if(condition):
      ##   A(i, j) = ...
    schedule: StageSchedule

  StageSchedule = ref object
  FunctionSchedule = ref object

  ParamKind = enum
    PkBuffer
    PkScalar
    PkTensor

  Parameter = ref object
    kind: ParamKind

  LuxNodeKind = enum
    # Expr
    lnkIntImm
    lnkFloatImm
    lnkAlgo
    lnkCall

    # Statement

  LuxNode = ref object
    case kind: LuxNodeKind
    of lnkIntImm:
      intVal: int
    of lnkFloatImm:
      floatVal: float
    of lnkAlgo:
      function: Function
    else:
      children: seq[LuxNode]

proc newImm(x: int): LuxNode =
  LuxNode(kind: lnkIntImm, intVal: x)

{.experimental:"callOperator".}

proc `()`(function: Function, params: varargs[LuxNode]): Call =
  Call(
    function: function,
    params: @params
  )

proc `:=`(call: Call, expression: LuxNode) =
  if call.function.isNil:
    new call.function

  let phaseId = call.function.stages.len

  call.function.stages.setLen(phaseId + 1)
  new call.function.stages[phaseId]
  shallowCopy(call.function.stages[phaseId].params, call.params)
  call.function.stages[phaseId].definition = expression

converter toLuxNode(call: Call): LuxNode =
  LuxNode(
    kind: lnkCall,
    children: @[
      LuxNode(kind: lnkAlgo, function: call.function),
      call.params[0],
      call.params[1]
    ]
  )

proc main() =

  var foo, bar: Function
  var x, y: LuxNode

  foo(x, y) := newImm 3
  bar(x, y) := foo(x, y)

main()