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
     ["[0]" maybe]
     ["[0]" try {"+" [Try]}]
     ["[0]" exception {"+" [exception:]}]
     ["<>" parser
      ["<[0]>" code {"+" [Parser]}]]]
    [data
     ["[0]" product]
     [collection
      ["[0]" list ("[1]\[0]" mix functor monoid)]
      ["[0]" array {"+" [Array]} ("[1]\[0]" functor mix)]]]
    [macro
     [syntax {"+" [syntax:]}]
     ["[0]" code]]
    [math
     [number
      ["n" nat]
      ["[0]" 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\[0]") (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 "[0]") 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 "[0]") ..functor
          (^open "[0]") ..mix
          (^open "[0]") ..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 "[0]") ..mix
          (^open "[0]") ..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]
  )