path: root/stdlib/source/lux/data/collection/row.lux

(.module:
  lux
  (lux (control [functor (#+ Functor)]
                [apply (#+ Apply)]
                [monad (#+ do Monad)]
                [equivalence (#+ Equivalence)]
                monoid
                fold
                ["p" parser])
       (data [maybe]
             (collection [list ("list/" Fold<List> Functor<List> Monoid<List>)]
                         [array ("array/" Functor<Array> Fold<Array>)])
             [bit]
             [product])
       [macro (#+ with-gensyms)]
       (macro [code]
              ["s" syntax (#+ syntax: Syntax)])
       ))

## [Utils]
(type: (Node a)
  (#Base (Array a))
  (#Hierarchy (Array (Node a))))

(type: (Base a) (Array a))
(type: (Hierarchy a) (Array (Node a)))

(type: Level Nat)

(type: Index Nat)

(def: branching-exponent
  Nat
  +5)

(def: root-level
  Level
  +0)

(do-template [<name> <op>]
  [(def: <name>
     (-> Level Level)
     (<op> branching-exponent))]

  [level-up   n/+]
  [level-down n/-]
  )

(def: full-node-size
  Nat
  (bit.left-shift branching-exponent +1))

(def: branch-idx-mask
  Nat
  (dec full-node-size))

(def: branch-idx
  (-> Index Index)
  (bit.and branch-idx-mask))

(def: (new-hierarchy _)
  (All [a] (-> Any (Hierarchy a)))
  (array.new full-node-size))

(def: (tail-off vec-size)
  (-> Nat Nat)
  (if (n/< full-node-size vec-size)
    +0
    (|> (dec vec-size)
        (bit.logical-right-shift branching-exponent)
        (bit.left-shift branching-exponent))))

(def: (new-path level tail)
  (All [a] (-> Level (Base a) (Node a)))
  (if (n/= +0 level)
    (#Base tail)
    (|> (new-hierarchy [])
        (array.write +0 (new-path (level-down level) tail))
        #Hierarchy)))

(def: (new-tail singleton)
  (All [a] (-> a (Base a)))
  (|> (array.new +1)
      (array.write +0 singleton)))

(def: (push-tail size level tail parent)
  (All [a] (-> Nat Level (Base a) (Hierarchy a) (Hierarchy a)))
  (let [sub-idx (branch-idx (bit.logical-right-shift level (dec size)))
        ## If we're currently on a bottom node
        sub-node (if (n/= branching-exponent level)
                   ## Just add the tail to it
                   (#Base tail)
                   ## Otherwise, check whether there's a vacant spot
                   (case (array.read sub-idx parent)
                     ## If so, set the path to the tail
                     #.None
                     (new-path (level-down level) tail)
                     ## If not, push the tail onto the sub-node.
                     (#.Some (#Hierarchy sub-node))
                     (#Hierarchy (push-tail size (level-down level) tail sub-node))

                     _
                     (undefined))
                   )]
    (|> (array.clone parent)
        (array.write sub-idx sub-node))))

(def: (expand-tail val tail)
  (All [a] (-> a (Base a) (Base a)))
  (let [tail-size (array.size tail)]
    (|> (array.new (inc tail-size))
        (array.copy tail-size +0 tail +0)
        (array.write tail-size val))))

(def: (put' level idx val hierarchy)
  (All [a] (-> Level Index a (Hierarchy a) (Hierarchy a)))
  (let [sub-idx (branch-idx (bit.logical-right-shift level idx))]
    (case (array.read sub-idx hierarchy)
      (#.Some (#Hierarchy sub-node))
      (|> (array.clone hierarchy)
          (array.write sub-idx (#Hierarchy (put' (level-down level) idx val sub-node))))

      (^multi (#.Some (#Base base))
              (n/= +0 (level-down level)))
      (|> (array.clone hierarchy)
          (array.write sub-idx (|> (array.clone base)
                                   (array.write (branch-idx idx) val)
                                   #Base)))

      _
      (undefined))))

(def: (pop-tail size level hierarchy)
  (All [a] (-> Nat Level (Hierarchy a) (Maybe (Hierarchy a))))
  (let [sub-idx (branch-idx (bit.logical-right-shift level (n/- +2 size)))]
    (cond (n/= +0 sub-idx)
          #.None

          (n/> branching-exponent level)
          (do maybe.Monad<Maybe>
            [base|hierarchy (array.read sub-idx hierarchy)
             sub (case base|hierarchy
                   (#Hierarchy sub)
                   (pop-tail size (level-down level) sub)

                   (#Base _)
                   (undefined))]
            (|> (array.clone hierarchy)
                (array.write sub-idx (#Hierarchy sub))
                #.Some))

          ## Else...
          (|> (array.clone hierarchy)
              (array.delete sub-idx)
              #.Some)
          )))

(def: (to-list' node)
  (All [a] (-> (Node a) (List a)))
  (case node
    (#Base base)
    (array.to-list base)
    
    (#Hierarchy hierarchy)
    (|> hierarchy
        array.to-list
        list.reverse
        (list/fold (function (_ sub acc) (list/compose (to-list' sub) acc))
                   #.Nil))))

## [Types]
(type: #export (Row a)
  {#level Level
   #size Nat
   #root (Hierarchy a)
   #tail (Base a)})

## [Exports]
(def: #export empty
  Row
  {#level (level-up root-level)
   #size +0
   #root (array.new full-node-size)
   #tail (array.new +0)})

(def: #export (size row)
  (All [a] (-> (Row a) Nat))
  (get@ #size row))

(def: #export (add val vec)
  (All [a] (-> a (Row a) (Row a)))
  ## Check if there is room in the tail.
  (let [vec-size (get@ #size vec)]
    (if (|> vec-size (n/- (tail-off vec-size)) (n/< full-node-size))
      ## If so, append to it.
      (|> vec
          (update@ #size inc)
          (update@ #tail (expand-tail val)))
      ## Otherwise, push tail into the tree
      ## --------------------------------------------------------
      ## Will the root experience an overflow with this addition?
      (|> (if (n/> (bit.left-shift (get@ #level vec) +1)
                   (bit.logical-right-shift branching-exponent vec-size))
            ## If so, a brand-new root must be established, that is
            ## 1-level taller.
            (|> vec
                (set@ #root (|> (: (Hierarchy ($ +0))
                                   (new-hierarchy []))
                                ## TODO: Remove version above once new-luxc becomes the standard compiler.
                                ## (new-hierarchy [])
                                (array.write +0 (#Hierarchy (get@ #root vec)))
                                (array.write +1 (new-path (get@ #level vec) (get@ #tail vec)))))
                (update@ #level level-up))
            ## Otherwise, just push the current tail onto the root.
            (|> vec
                (update@ #root (push-tail vec-size (get@ #level vec) (get@ #tail vec)))))
          ## Finally, update the size of the row and grow a new
          ## tail with the new element as it's sole member.
          (update@ #size inc)
          (set@ #tail (new-tail val)))
      )))

(def: (base-for idx vec)
  (All [a] (-> Index (Row a) (Maybe (Base a))))
  (let [vec-size (get@ #size vec)]
    (if (and (n/>= +0 idx)
             (n/< vec-size idx))
      (if (n/>= (tail-off vec-size) idx)
        (#.Some (get@ #tail vec))
        (loop [level (get@ #level vec)
               hierarchy (get@ #root vec)]
          (case [(n/> branching-exponent level)
                 (array.read (branch-idx (bit.logical-right-shift level idx)) hierarchy)]
            [true (#.Some (#Hierarchy sub))]
            (recur (level-down level) sub)

            [false (#.Some (#Base base))]
            (#.Some base)

            [_ #.None]
            #.None

            _
            (error! "Incorrect row structure."))))
      #.None)))

(def: #export (nth idx vec)
  (All [a] (-> Nat (Row a) (Maybe a)))
  (do maybe.Monad<Maybe>
    [base (base-for idx vec)]
    (array.read (branch-idx idx) base)))

(def: #export (put idx val vec)
  (All [a] (-> Nat a (Row a) (Row a)))
  (let [vec-size (get@ #size vec)]
    (if (and (n/>= +0 idx)
             (n/< vec-size idx))
      (if (n/>= (tail-off vec-size) idx)
        (|> vec
            ## (update@ #tail (|>> array.clone (array.write (branch-idx idx) val)))
            ## TODO: Remove once new-luxc becomes the standard compiler.
            (update@ #tail (: (-> (Base ($ +0)) (Base ($ +0)))
                              (|>> array.clone (array.write (branch-idx idx) val))))
            )
        (|> vec
            (update@ #root (put' (get@ #level vec) idx val))))
      vec)))

(def: #export (update idx f vec)
  (All [a] (-> Nat (-> a a) (Row a) (Row a)))
  (case (nth idx vec)
    (#.Some val)
    (put idx (f val) vec)

    #.None
    vec))

(def: #export (pop vec)
  (All [a] (-> (Row a) (Row a)))
  (case (get@ #size vec)
    +0
    empty

    +1
    empty

    vec-size
    (if (|> vec-size (n/- (tail-off vec-size)) (n/> +1))
      (let [old-tail (get@ #tail vec)
            new-tail-size (dec (array.size old-tail))]
        (|> vec
            (update@ #size dec)
            (set@ #tail (|> (array.new new-tail-size)
                            (array.copy new-tail-size +0 old-tail +0)))))
      (maybe.assume
       (do maybe.Monad<Maybe>
         [new-tail (base-for (n/- +2 vec-size) vec)
          #let [[level' root'] (let [init-level (get@ #level vec)]
                                 (loop [level init-level
                                        root (maybe.default (new-hierarchy [])
                                                            (pop-tail vec-size init-level (get@ #root vec)))]
                                   (if (n/> branching-exponent level)
                                     (case [(array.read +1 root) (array.read +0 root)]
                                       [#.None (#.Some (#Hierarchy sub-node))]
                                       (recur (level-down level) sub-node)

                                       ## [#.None (#.Some (#Base _))]
                                       ## (undefined)

                                       _
                                       [level root])
                                     [level root])))]]
         (wrap (|> vec
                   (update@ #size dec)
                   (set@ #level level')
                   (set@ #root root')
                   (set@ #tail new-tail))))))
    ))

(def: #export (to-list vec)
  (All [a] (-> (Row a) (List a)))
  (list/compose (to-list' (#Hierarchy (get@ #root vec)))
                (to-list' (#Base (get@ #tail vec)))))

(def: #export (from-list list)
  (All [a] (-> (List a) (Row a)))
  (list/fold add
             empty
             list))

(def: #export (member? a/Equivalence vec val)
  (All [a] (-> (Equivalence a) (Row a) a Bool))
  (list.member? a/Equivalence (to-list vec) val))

(def: #export empty?
  (All [a] (-> (Row a) Bool))
  (|>> (get@ #size) (n/= +0)))

## [Syntax]
(syntax: #export (row {elems (p.some s.any)})
  {#.doc (doc "Row literals."
              (row 10 20 30 40))}
  (wrap (list (` (from-list (list (~+ elems)))))))

## [Structures]
(structure: #export (Equivalence<Node> Equivalence<a>) (All [a] (-> (Equivalence a) (Equivalence (Node a))))
  (def: (= v1 v2)
    (case [v1 v2]
      [(#Base b1) (#Base b2)]
      (:: (array.Equivalence<Array> Equivalence<a>) = b1 b2)
      
      [(#Hierarchy h1) (#Hierarchy h2)]
      (:: (array.Equivalence<Array> (Equivalence<Node> Equivalence<a>)) = h1 h2)

      _
      false)))

(structure: #export (Equivalence<Row> Equivalence<a>) (All [a] (-> (Equivalence a) (Equivalence (Row a))))
  (def: (= v1 v2)
    (and (n/= (get@ #size v1) (get@ #size v2))
         (let [(^open "Node/") (Equivalence<Node> Equivalence<a>)]
           (and (Node/= (#Base (get@ #tail v1))
                        (#Base (get@ #tail v2)))
                (Node/= (#Hierarchy (get@ #root v1))
                        (#Hierarchy (get@ #root v2))))))))

(structure: _ (Fold Node)
  (def: (fold f init xs)
    (case xs
      (#Base base)
      (array/fold f init base)
      
      (#Hierarchy hierarchy)
      (array/fold (function (_ node init') (fold f init' node))
                  init
                  hierarchy))
    ))

(structure: #export _ (Fold Row)
  (def: (fold f init xs)
    (let [(^open) Fold<Node>]
      (fold f
            (fold f
                  init
                  (#Hierarchy (get@ #root xs)))
            (#Base (get@ #tail xs))))
    ))

(structure: #export Monoid<Row> (All [a] (Monoid (Row a)))
  (def: identity empty)
  (def: (compose xs ys)
    (list/fold add xs (to-list ys))))

(structure: _ (Functor Node)
  (def: (map f xs)
    (case xs
      (#Base base)
      (#Base (array/map f base))
      
      (#Hierarchy hierarchy)
      (#Hierarchy (array/map (map f) hierarchy)))
    ))

(structure: #export _ (Functor Row)
  (def: (map f xs)
    {#level (get@ #level xs)
     #size (get@ #size xs)
     #root (|> xs (get@ #root) (array/map (:: Functor<Node> map f)))
     #tail (|> xs (get@ #tail) (array/map f))
     }))

(structure: #export _ (Apply Row)
  (def: functor Functor<Row>)

  (def: (apply ff fa)
    (let [(^open) Functor<Row>
          (^open) Fold<Row>
          (^open) Monoid<Row>
          results (map (function (_ f) (map f fa))
                       ff)]
      (fold compose identity results))))

(structure: #export _ (Monad Row)
  (def: functor Functor<Row>)

  (def: wrap (|>> row))

  (def: join
    (let [(^open) Fold<Row>
          (^open) Monoid<Row>]
      (fold (function (_ post pre) (compose pre post)) identity))))

(def: #export (reverse xs)
  (All [a] (-> (Row a) (Row a)))
  (let [(^open) Fold<Row>
        (^open) Monoid<Row>]
    (fold add identity xs)))