(.module:
  [lux #*
   [abstract
    functor
    [monad (#+ do Monad)]
    equivalence
    fold]
   [control
    ["p" parser
     ["s" code (#+ Parser)]]]
   [data
    [collection
     ["." list ("#@." monad fold)]]]
   ["." macro
    ["." code]
    [syntax (#+ syntax:)]]])

(type: #export (Tree a)
  {#value a
   #children (List (Tree a))})

(def: #export (flatten tree)
  (All [a] (-> (Tree a) (List a)))
  (#.Cons (get@ #value tree)
          (list@join (list@map flatten (get@ #children tree)))))

(def: #export (leaf value)
  (All [a] (-> a (Tree a)))
  {#value value
   #children (list)})

(def: #export (branch value children)
  (All [a] (-> a (List (Tree a)) (Tree a)))
  {#value value
   #children children})

(type: #rec Tree-Code
  [Code (List Tree-Code)])

(def: tree^
  (Parser Tree-Code)
  (|> (|>> p.some s.record (p.and s.any))
      p.rec
      p.some
      s.record
      (p.and s.any)
      s.tuple))

(syntax: #export (tree {root tree^})
  {#.doc (doc "Tree literals."
              (tree Int [+10 {+20 {}
                              +30 {}
                              +40 {}}]))}
  (wrap (list (` (~ (loop [[value children] root]
                      (` {#value (~ value)
                          #children (list (~+ (list@map recur children)))})))))))

(structure: #export (equivalence Equivalence<a>) (All [a] (-> (Equivalence a) (Equivalence (Tree a))))
  (def: (= tx ty)
    (and (:: Equivalence<a> = (get@ #value tx) (get@ #value ty))
         (:: (list.equivalence (equivalence Equivalence<a>)) = (get@ #children tx) (get@ #children ty)))))

(structure: #export functor (Functor Tree)
  (def: (map f fa)
    {#value (f (get@ #value fa))
     #children (list@map (map f)
                         (get@ #children fa))}))

(structure: #export fold (Fold Tree)
  (def: (fold f init tree)
    (list@fold (function (_ tree' init') (fold f init' tree'))
               (f (get@ #value tree)
                  init)
               (get@ #children tree))))