ast.nim

when not defined(js):
  import slice
  import ../stream
  import ../collections/[seqstack]
  import ../monad/[identity, optional]
  import ../optics/[lens]
  import ../utils/[nimnodes]

  import std/[macros, sugar]



  export nimnodes, stream



  type
    NimNodeStep* = SliceStep[NimNodeIndex]

    PreOrderStep* = object
      ## Since 0.3.0.
      stack: SeqStack[NimNode]
      current: Optional[NimNode]



  func indexes* (n: NimNode): Stream[NimNodeStep, NimNodeIndex] =
    slice(n.low(), n.high()).items()


  func indexesReverse* (n: NimNode): Stream[NimNodeStep, NimNodeIndex] =
    ## Since 0.3.0.
    slice(n.low(), n.high()).itemsReverse()


  func children* (n: NimNode): Stream[NimNodeStep, NimNode] =
    n.indexes().map(i => n[i])


  func childrenReverse* (n: NimNode): Stream[NimNodeStep, NimNode] =
    ## Since 0.3.0.
    n.indexesReverse().map(i => n[i])


  func pairs* (
    n: NimNode
  ): Stream[NimNodeStep, tuple[index: NimNodeIndex; node: NimNode]] =
    n.indexes().map(i => (i, n[i]))



  func preOrderStep* (
    stack: SeqStack[NimNode];
    current: Optional[NimNode]
  ): PreOrderStep =
    ## Since 0.3.0.
    PreOrderStep(stack: stack, current: current)



  func current (X: typedesc[PreOrderStep]): Lens[X, NimNode] =
    lens(
      (self: X) => self.current,
      (self: X, current: Optional[NimNode]) => preOrderStep(self.stack, current)
    )


  func stack (X: typedesc[PreOrderStep]): Lens[X, SeqStack[NimNode]] =
    lens(
      (self: X) => self.stack,
      (self: X, stack: SeqStack[NimNode]) => preOrderStep(stack, self.current)
    )



  func initPreOrder* (root: NimNode): PreOrderStep =
    ## Since 0.3.0.
    preOrderStep(seqStack[NimNode](), root.toSome())


  func hasMore* (self: PreOrderStep): bool =
    ## Since 0.3.0.
    self.current.isSome()


  func generate* (self: PreOrderStep): NimNode {.
    raises: [Exception, UnboxError]
  .} =
    ## Since 0.3.0.
    self.current.unbox()


  func next* (self: PreOrderStep): PreOrderStep =
    ## Since 0.3.0.
    self
      .generate()
      .childrenReverse()
      .reduce(push[NimNode], self.stack)
      .pop()
      .apply(popResult => preOrderStep(popResult.stack, popResult.popped))



  func traversePreOrder* (root: NimNode): Stream[PreOrderStep, NimNode] =
    ##[
      Traverse the tree starting at `root` in pre-order.

      Since 0.3.0.
    ]##
    hasMore
      .looped(next)
      .generating(generate)
      .startingAt(() => root.initPreOrder())



  when isMainModule:
    import ../utils/[call, convert, ignore, partialprocs, unit]

    import std/[os, unittest]



    proc main () =
      suite currentSourcePath().splitFile().name:
        test """Counting the number of children of a NimNode "n" through a stream should return "n.len()".""":
          template doTest (n: NimNode): proc () =
            (
              proc () =
                const
                  actual = n.children().count(Natural)
                  expected = n.len()

                check:
                  actual == expected
            )


          for t in [
            doTest(newEmptyNode()),
            doTest(
              quote do:
                var a = nil
                let b = a

                echo(b)
            )
          ]:
            t.call()



        test """All the items generated by "n.pairs()" should be equal to the index and child NimNode of "n" at that index.""":
          template doTest (n: NimNode): proc () =
            (
              proc () =
                const verified =
                  n.pairs().all(
                    (pair: (NimNodeIndex, NimNode)) => pair[1] == n[pair[0]]
                  )

                check:
                  verified
            )


          for t in [
            doTest(newEmptyNode()),
            doTest("abc".newStrLitNode()),
            doTest(newProc())
          ]:
            t.call()



        test """"root.traversePreOrder()" should yield itself and its descendants in depth first pre-order.""":
          # "Seq[NimNode]" is not a valid type for "const" symbols.

          func collectPreOrder (root: NimNode): seq[NimNode] =
            var stack = @[root]
            result = @[]

            while stack.len() > 0:
              let current = stack.pop()

              result.add(current)
              for i in current.low() .. current.high():
                let j = current.high() - i

                stack.add(current[j])


          proc doTest (root: NimNode) =
            let
              actual =
                root
                  .traversePreOrder()
                  .reduce(
                    partial(?:seq[NimNode] & ?:NimNode),
                    seq[NimNode](@[])
                  )
              expected = root.collectPreOrder()

            doAssert(actual == expected)



          static:
            doTest(newEmptyNode())
            doTest("abc".newLit())
            doTest("doTest".newCall(0.newLit(), 0.0.newLit()))
            doTest(
              quote do:
                var a = "abc"
                let b = a.f(a & a, g(a))
            )



    main()