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

... https://hypirion.com/musings/understanding-persistent-vector-pt-1
... https://hypirion.com/musings/understanding-persistent-vector-pt-2
... https://hypirion.com/musings/understanding-persistent-vector-pt-3
(.module:
  [library
   [lux (#- list)
    ["@" target]
    [abstract
     [functor (#+ Functor)]
     [apply (#+ Apply)]
     [monad (#+ Monad do)]
     [equivalence (#+ Equivalence)]
     [monoid (#+ Monoid)]
     [mix (#+ Mix)]
     [predicate (#+ Predicate)]]
    [control
     ["." maybe]
     ["." try (#+ Try)]
     ["." exception (#+ exception:)]
     ["<>" parser
      ["<.>" code (#+ Parser)]]]
    [data
     ["." product]
     [collection
      ["." list ("#\." mix functor monoid)]
      ["." array (#+ Array) ("#\." functor mix)]]]
    [macro
     [syntax (#+ syntax:)]
     ["." code]]
    [math
     [number
      ["n" nat]
      ["." i64]]]]])

(type: (Node a)
  (Variant
   (#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)

(template [<name> <op>]
  [(def: <name>
     (-> Level Level)
     (<op> branching_exponent))]

  [level_up   n.+]
  [level_down n.-]
  )

(def: full_node_size
  Nat
  (i64.left_shifted branching_exponent 1))

(def: branch_idx_mask
  Nat
  (-- full_node_size))

(def: branch_idx
  (-> Index Index)
  (i64.and branch_idx_mask))

(def: (empty_hierarchy _)
  (All (_ a) (-> Any (Hierarchy a)))
  (array.empty ..full_node_size))

(def: (tail_off row_size)
  (-> Nat Nat)
  (if (n.< full_node_size row_size)
    0
    (|> (-- row_size)
        (i64.right_shifted branching_exponent)
        (i64.left_shifted branching_exponent))))

(def: (path level tail)
  (All (_ a) (-> Level (Base a) (Node a)))
  (if (n.= 0 level)
    (#Base tail)
    (|> (empty_hierarchy [])
        (array.write! 0 (path (level_down level) tail))
        #Hierarchy)))

(def: (tail singleton)
  (All (_ a) (-> a (Base a)))
  (|> (array.empty 1)
      (array.write! 0 singleton)))

(def: (with_tail size level tail parent)
  (All (_ a) (-> Nat Level (Base a) (Hierarchy a) (Hierarchy a)))
  (let [sub_idx (branch_idx (i64.right_shifted level (-- 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
                     (..path (level_down level) tail)
                     ... If not, push the tail onto the sub_node.
                     (#.Some (#Hierarchy sub_node))
                     (#Hierarchy (with_tail size (level_down level) tail sub_node))

                     _
                     (undefined))
                   )]
    (|> (array.clone parent)
        (array.write! sub_idx sub_node))))

(def: (expanded_tail val tail)
  (All (_ a) (-> a (Base a) (Base a)))
  (let [tail_size (array.size tail)]
    (|> (array.empty (++ tail_size))
        (array.copy! tail_size 0 tail 0)
        (array.write! tail_size val))))

(def: (hierarchy\has level idx val hierarchy)
  (All (_ a) (-> Level Index a (Hierarchy a) (Hierarchy a)))
  (let [sub_idx (branch_idx (i64.right_shifted level idx))]
    (case (array.read! sub_idx hierarchy)
      (#.Some (#Hierarchy sub_node))
      (|> (array.clone hierarchy)
          (array.write! sub_idx (#Hierarchy (hierarchy\has (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: (without_tail size level hierarchy)
  (All (_ a) (-> Nat Level (Hierarchy a) (Maybe (Hierarchy a))))
  (let [sub_idx (branch_idx (i64.right_shifted level (n.- 2 size)))]
    (cond (n.= 0 sub_idx)
          #.None

          (n.> branching_exponent level)
          (do maybe.monad
            [base|hierarchy (array.read! sub_idx hierarchy)
             sub (case base|hierarchy
                   (#Hierarchy sub)
                   (without_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: (node\list node)
  (All (_ a) (-> (Node a) (List a)))
  (case node
    (#Base base)
    (array.list #.None base)
    
    (#Hierarchy hierarchy)
    (|> hierarchy
        (array.list #.None)
        list.reversed
        (list\mix (function (_ sub acc)
                    (list\composite (node\list sub) acc))
                  #.End))))

(type: .public (Row a)
  (Record
   [#level Level
    #size Nat
    #root (Hierarchy a)
    #tail (Base a)]))

(def: .public empty
  Row
  [#level (level_up root_level)
   #size 0
   #root (empty_hierarchy [])
   #tail (array.empty 0)])

(def: .public (size row)
  (All (_ a) (-> (Row a) Nat))
  (value@ #size row))

(def: .public (suffix val row)
  (All (_ a) (-> a (Row a) (Row a)))
  ... Check if there is room in the tail.
  (let [row_size (value@ #size row)]
    (if (|> row_size (n.- (tail_off row_size)) (n.< full_node_size))
      ... If so, append to it.
      (|> row
          (revised@ #size ++)
          (revised@ #tail (..expanded_tail val)))
      ... Otherwise, push tail into the tree
      ... --------------------------------------------------------
      ... Will the root experience an overflow with this addition?
      (|> (if (n.> (i64.left_shifted (value@ #level row) 1)
                   (i64.right_shifted branching_exponent row_size))
            ... If so, a brand-new root must be established, that is
            ... 1-level taller.
            (|> row
                (with@ #root (|> (for {@.old
                                       (: (Hierarchy (:parameter 0))
                                          (empty_hierarchy []))}
                                      (empty_hierarchy []))
                                 (array.write! 0 (#Hierarchy (value@ #root row)))
                                 (array.write! 1 (..path (value@ #level row) (value@ #tail row)))))
                (revised@ #level level_up))
            ... Otherwise, just push the current tail onto the root.
            (|> row
                (revised@ #root (..with_tail row_size (value@ #level row) (value@ #tail row)))))
          ... Finally, update the size of the row and grow a new
          ... tail with the new element as it's sole member.
          (revised@ #size ++)
          (with@ #tail (..tail val)))
      )))

(exception: incorrect_row_structure)

(exception: .public [a] (index_out_of_bounds {row (Row a)} {index Nat})
  (exception.report ["Size" (\ n.decimal encoded (value@ #size row))]
                    ["Index" (\ n.decimal encoded index)]))

(exception: base_was_not_found)

(def: .public (within_bounds? row idx)
  (All (_ a) (-> (Row a) Nat Bit))
  (n.< (value@ #size row) idx))

(def: (base_for idx row)
  (All (_ a) (-> Index (Row a) (Try (Base a))))
  (if (within_bounds? row idx)
    (if (n.< (tail_off (value@ #size row)) idx)
      (loop [level (value@ #level row)
             hierarchy (value@ #root row)]
        (case [(n.> branching_exponent level)
               (array.read! (branch_idx (i64.right_shifted level idx)) hierarchy)]
          [#1 (#.Some (#Hierarchy sub))]
          (recur (level_down level) sub)

          [#0 (#.Some (#Base base))]
          (#try.Success base)

          [_ #.None]
          (exception.except ..base_was_not_found [])

          _
          (exception.except ..incorrect_row_structure [])))
      (#try.Success (value@ #tail row)))
    (exception.except ..index_out_of_bounds [row idx])))

(def: .public (item idx row)
  (All (_ a) (-> Nat (Row a) (Try a)))
  (do try.monad
    [base (base_for idx row)]
    (case (array.read! (branch_idx idx) base)
      (#.Some value)
      (#try.Success value)
      
      #.None
      (exception.except ..incorrect_row_structure []))))

(def: .public (has idx val row)
  (All (_ a) (-> Nat a (Row a) (Try (Row a))))
  (let [row_size (value@ #size row)]
    (if (within_bounds? row idx)
      (#try.Success (if (n.< (tail_off row_size) idx)
                      (revised@ #root (hierarchy\has (value@ #level row) idx val)
                                row)
                      (revised@ #tail (for {@.old
                                            (: (-> (Base (:parameter 0)) (Base (:parameter 0)))
                                               (|>> array.clone (array.write! (branch_idx idx) val)))}
                                           (|>> array.clone (array.write! (branch_idx idx) val)))
                                row)))
      (exception.except ..index_out_of_bounds [row idx]))))

(def: .public (revised idx f row)
  (All (_ a) (-> Nat (-> a a) (Row a) (Try (Row a))))
  (do try.monad
    [val (..item idx row)]
    (..has idx (f val) row)))

(def: .public (prefix row)
  (All (_ a) (-> (Row a) (Row a)))
  (case (value@ #size row)
    0
    empty

    1
    empty

    row_size
    (if (|> row_size (n.- (tail_off row_size)) (n.> 1))
      (let [old_tail (value@ #tail row)
            new_tail_size (-- (array.size old_tail))]
        (|> row
            (revised@ #size --)
            (with@ #tail (|> (array.empty new_tail_size)
                             (array.copy! new_tail_size 0 old_tail 0)))))
      (maybe.trusted
       (do maybe.monad
         [new_tail (base_for (n.- 2 row_size) row)
          .let [[level' root'] (let [init_level (value@ #level row)]
                                 (loop [level init_level
                                        root (maybe.else (empty_hierarchy [])
                                                         (without_tail row_size init_level (value@ #root row)))]
                                   (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])))]]
         (in (|> row
                 (revised@ #size --)
                 (with@ #level level')
                 (with@ #root root')
                 (with@ #tail new_tail))))))
    ))

(def: .public (list row)
  (All (_ a) (-> (Row a) (List a)))
  (list\composite (node\list (#Hierarchy (value@ #root row)))
                  (node\list (#Base (value@ #tail row)))))

(def: .public of_list
  (All (_ a) (-> (List a) (Row a)))
  (list\mix ..suffix ..empty))

(def: .public (member? equivalence row val)
  (All (_ a) (-> (Equivalence a) (Row a) a Bit))
  (list.member? equivalence (list row) val))

(def: .public empty?
  (All (_ a) (-> (Row a) Bit))
  (|>> (value@ #size) (n.= 0)))

(syntax: .public (row [elems (<>.some <code>.any)])
  (in (.list (` (..of_list (.list (~+ elems)))))))

(implementation: (node_equivalence Equivalence<a>)
  (All (_ a) (-> (Equivalence a) (Equivalence (Node a))))
  
  (def: (= v1 v2)
    (case [v1 v2]
      [(#Base b1) (#Base b2)]
      (\ (array.equivalence Equivalence<a>) = b1 b2)
      
      [(#Hierarchy h1) (#Hierarchy h2)]
      (\ (array.equivalence (node_equivalence Equivalence<a>)) = h1 h2)

      _
      #0)))

(implementation: .public (equivalence Equivalence<a>)
  (All (_ a) (-> (Equivalence a) (Equivalence (Row a))))
  
  (def: (= v1 v2)
    (and (n.= (value@ #size v1) (value@ #size v2))
         (let [(^open "node\.") (node_equivalence Equivalence<a>)]
           (and (node\= (#Base (value@ #tail v1))
                        (#Base (value@ #tail v2)))
                (node\= (#Hierarchy (value@ #root v1))
                        (#Hierarchy (value@ #root v2))))))))

(implementation: node_mix
  (Mix Node)
  
  (def: (mix f init xs)
    (case xs
      (#Base base)
      (array\mix f init base)
      
      (#Hierarchy hierarchy)
      (array\mix (function (_ node init') (mix f init' node))
                 init
                 hierarchy))))

(implementation: .public mix
  (Mix Row)
  
  (def: (mix f init xs)
    (let [(^open ".") node_mix]
      (mix f
           (mix f
                init
                (#Hierarchy (value@ #root xs)))
           (#Base (value@ #tail xs))))))

(implementation: .public monoid
  (All (_ a) (Monoid (Row a)))
  
  (def: identity ..empty)
  
  (def: (composite xs ys)
    (list\mix suffix xs (..list ys))))

(implementation: node_functor
  (Functor Node)
  
  (def: (each f xs)
    (case xs
      (#Base base)
      (#Base (array\each f base))
      
      (#Hierarchy hierarchy)
      (#Hierarchy (array\each (each f) hierarchy)))))

(implementation: .public functor
  (Functor Row)
  
  (def: (each f xs)
    [#level (value@ #level xs)
     #size (value@ #size xs)
     #root (|> xs (value@ #root) (array\each (\ node_functor each f)))
     #tail (|> xs (value@ #tail) (array\each f))]))

(implementation: .public apply
  (Apply Row)
  
  (def: &functor ..functor)

  (def: (on fa ff)
    (let [(^open ".") ..functor
          (^open ".") ..mix
          (^open ".") ..monoid
          results (each (function (_ f) (each f fa))
                        ff)]
      (mix composite identity results))))

(implementation: .public monad
  (Monad Row)
  
  (def: &functor ..functor)

  (def: in
    (|>> row))

  (def: conjoint
    (let [(^open ".") ..mix
          (^open ".") ..monoid]
      (mix (function (_ post pre) (composite pre post)) identity))))

(def: .public reversed
  (All (_ a) (-> (Row a) (Row a)))
  (|>> ..list
       list.reversed
       (list\mix suffix ..empty)))

(template [<name> <array> <init> <op>]
  [(def: .public <name>
     (All (_ a)
       (-> (Predicate a) (Row a) Bit))
     (let [help (: (All (_ a)
                     (-> (Predicate a) (Node a) Bit))
                   (function (help predicate node)
                     (case node
                       (#Base base)
                       (<array> predicate base)

                       (#Hierarchy hierarchy)
                       (<array> (help predicate) hierarchy))))]
       (function (<name> predicate row)
         (let [(^slots [#root #tail]) row]
           (<op> (help predicate (#Hierarchy root))
                 (help predicate (#Base tail)))))))]

  [every? array.every? #1 and]
  [any?   array.any?   #0 or]
  )