path: root/stdlib/source/library/lux/data/collection/sequence.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
(.require
 [library
  [lux (.except list has revised only)
   [abstract
    [functor (.only Functor)]
    [apply (.only Apply)]
    [monad (.only Monad do)]
    [equivalence (.only Equivalence)]
    [monoid (.only Monoid)]
    [mix (.only Mix)]]
   [control
    ["<>" parser]
    ["[0]" maybe (.use "[1]#[0]" functor)]
    ["[0]" try (.only Try)]
    ["[0]" exception (.only Exception)]
    [function
     [predicate (.only Predicate)]]]
   [data
    ["[0]" product]
    [collection
     ["[0]" list (.use "[1]#[0]" mix functor monoid)]
     ["[0]" array
      ["[1]" \\unsafe (.only Array)]]]]
   [math
    [number
     ["n" nat]
     ["[0]" i64]]]
   [meta
    ["@" target]
    ["[0]" code (.only)
     ["<[1]>" \\parser (.only Parser)]]
    [macro
     [syntax (.only syntax)]
     ["^" pattern]]]]])

(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)

(with_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 sequence_size)
  (-> Nat Nat)
  (if (n.< full_node_size sequence_size)
    0
    (|> (-- sequence_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.has! 0 (path (level_down level) tail))
        {#Hierarchy})))

(def (tail singleton)
  (All (_ a) (-> a (Base a)))
  (|> (array.empty 1)
      (array.has! 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
                   (if (array.lacks? sub_idx parent)
                     ... If so, set the path to the tail
                     (..path (level_down level) tail)
                     (when (array.item sub_idx parent)
                       ... If not, push the tail onto the sub_node.
                       {#Hierarchy sub_node}
                       {#Hierarchy (with_tail size (level_down level) tail sub_node)}

                       _
                       (undefined))))]
    (|> (array.clone parent)
        (array.has! 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.has! 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))]
    (when (array.item sub_idx hierarchy)
      {#Hierarchy sub_node}
      (|> (array.clone hierarchy)
          (array.has! sub_idx {#Hierarchy (hierarchy#has (level_down level) idx val sub_node)}))

      (^.multi {#Base base}
               (n.= 0 (level_down level)))
      (|> (array.clone hierarchy)
          (array.has! sub_idx (|> (array.clone base)
                                  (array.has! (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)
          (if (array.lacks? sub_idx hierarchy)
            {.#None}
            (maybe#each (function (_ sub)
                          (|> (array.clone hierarchy)
                              (array.has! sub_idx {#Hierarchy sub})))
                        (when (array.item sub_idx hierarchy)
                          {#Hierarchy sub}
                          (without_tail size (level_down level) sub)

                          {#Base _}
                          (undefined))))

          ... Else...
          (|> (array.clone hierarchy)
              (array.lacks! sub_idx)
              {.#Some})
          )))

(def (node#list node)
  (All (_ a) (-> (Node a) (List a)))
  (when 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 (Sequence a)
  (Record
   [#level Level
    #size Nat
    #root (Hierarchy a)
    #tail (Base a)]))

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

(def .public (size sequence)
  (All (_ a) (-> (Sequence a) Nat))
  (the #size sequence))

(def .public (suffix val sequence)
  (All (_ a) (-> a (Sequence a) (Sequence a)))
  ... Check if there is room in the tail.
  (let [sequence_size (the #size sequence)]
    (if (|> sequence_size (n.- (tail_off sequence_size)) (n.< full_node_size))
      ... If so, append to it.
      (|> sequence
          (.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 (the #level sequence) 1)
                   (i64.right_shifted branching_exponent sequence_size))
            ... If so, a brand-new root must be established, that is
            ... 1-level taller.
            (|> sequence
                (.has #root (|> (`` (is (Hierarchy (,, (type_of val)))
                                        (empty_hierarchy [])))
                                (array.has! 0 {#Hierarchy (the #root sequence)})
                                (array.has! 1 (..path (the #level sequence) (the #tail sequence)))))
                (.revised #level level_up))
            ... Otherwise, just push the current tail onto the root.
            (|> sequence
                (.revised #root (..with_tail sequence_size (the #level sequence) (the #tail sequence)))))
          ... Finally, update the size of the sequence and grow a new
          ... tail with the new element as it's sole member.
          (.revised #size ++)
          (.has #tail (..tail val)))
      )))

(exception.def incorrect_sequence_structure)

(exception.def .public (index_out_of_bounds [sequence index])
  (All (_ a) (Exception [(Sequence a) Nat]))
  (exception.report
   (.list ["Size" (at n.decimal encoded (the #size sequence))]
          ["Index" (at n.decimal encoded index)])))

(exception.def base_was_not_found)

(def .public (within_bounds? sequence idx)
  (All (_ a) (-> (Sequence a) Nat Bit))
  (n.< (the #size sequence) idx))

(def (base_for idx sequence)
  (All (_ a) (-> Index (Sequence a) (Try (Base a))))
  (if (within_bounds? sequence idx)
    (if (n.< (tail_off (the #size sequence)) idx)
      (loop (again [level (the #level sequence)
                    hierarchy (the #root sequence)])
        (let [index (branch_idx (i64.right_shifted level idx))]
          (if (array.lacks? index hierarchy)
            (exception.except ..base_was_not_found [])
            (when [(n.> branching_exponent level)
                   (array.item index hierarchy)]
              [.true {#Hierarchy sub}]
              (again (level_down level) sub)

              [.false {#Base base}]
              {try.#Success base}

              _
              (exception.except ..incorrect_sequence_structure [])))))
      {try.#Success (the #tail sequence)})
    (exception.except ..index_out_of_bounds [sequence idx])))

(def .public (item idx sequence)
  (All (_ a) (-> Nat (Sequence a) (Try a)))
  (do try.monad
    [base (base_for idx sequence)
     .let [index (branch_idx idx)]]
    (if (array.lacks? index base)
      (exception.except ..incorrect_sequence_structure [])
      {try.#Success (array.item index base)})))

(def .public (has idx val sequence)
  (All (_ a) (-> Nat a (Sequence a) (Try (Sequence a))))
  (let [sequence_size (the #size sequence)]
    (if (within_bounds? sequence idx)
      {try.#Success (if (n.< (tail_off sequence_size) idx)
                      (.revised #root (hierarchy#has (the #level sequence) idx val)
                                sequence)
                      (.revised #tail (`` (is (-> (Base (,, (type_of val)))
                                                  (Base (,, (type_of val))))
                                              (|>> array.clone (array.has! (branch_idx idx) val))))
                                sequence))}
      (exception.except ..index_out_of_bounds [sequence idx]))))

(def .public (revised idx revision it)
  (All (_ a) (-> Nat (-> a a) (Sequence a) (Try (Sequence a))))
  (do try.monad
    [val (..item idx it)]
    (..has idx (revision val) it)))

(def .public (prefix sequence)
  (All (_ a) (-> (Sequence a) (Sequence a)))
  (when (the #size sequence)
    0
    empty

    1
    empty

    sequence_size
    (if (|> sequence_size (n.- (tail_off sequence_size)) (n.> 1))
      (let [old_tail (the #tail sequence)
            new_tail_size (-- (array.size old_tail))]
        (|> sequence
            (.revised #size --)
            (.has #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 sequence_size) sequence)
          .let [[level' root'] (let [init_level (the #level sequence)]
                                 (loop (again [level init_level
                                               root (maybe.else (empty_hierarchy [])
                                                                (without_tail sequence_size init_level (the #root sequence)))])
                                   (with_expansions [<else> [level root]]
                                     (if (n.> branching_exponent level)
                                       (if (array.lacks? 1 root)
                                         (when (array.item 0 root)
                                           {#Hierarchy sub_node}
                                           (again (level_down level) sub_node)

                                           ... {#Base _}
                                           ... (undefined)

                                           _
                                           <else>)
                                         <else>)
                                       <else>))))]]
         (in (|> sequence
                 (.revised #size --)
                 (.has #level level')
                 (.has #root root')
                 (.has #tail new_tail))))))
    ))

(def .public (list sequence)
  (All (_ a) (-> (Sequence a) (List a)))
  (list#composite (node#list {#Hierarchy (the #root sequence)})
                  (node#list {#Base (the #tail sequence)})))

(def .public of_list
  (All (_ a) (-> (List a) (Sequence a)))
  (list#mix ..suffix ..empty))

(def .public (member? equivalence sequence val)
  (All (_ a) (-> (Equivalence a) (Sequence a) a Bit))
  (list.member? equivalence (list sequence) val))

(def .public empty?
  (All (_ a) (-> (Sequence a) Bit))
  (|>> (the #size) (n.= 0)))

(def .public sequence
  (syntax (_ [elems (<>.some <code>.any)])
    (in (.list (` (..of_list (.list (,* elems))))))))

(def (node_equivalence //#=)
  (All (_ a) (-> (Equivalence a) (Equivalence (Node a))))
  (implementation
   (def (= v1 v2)
     (when [v1 v2]
       [{#Base b1} {#Base b2}]
       (array.= //#= b1 b2)
       
       [{#Hierarchy h1} {#Hierarchy h2}]
       (array.= (node_equivalence //#=) h1 h2)

       _
       false))))

(def .public (equivalence //#=)
  (All (_ a) (-> (Equivalence a) (Equivalence (Sequence a))))
  (implementation
   (def (= v1 v2)
     (and (n.= (the #size v1) (the #size v2))
          (let [(open "node#[0]") (node_equivalence //#=)]
            (and (node#= {#Base (the #tail v1)}
                         {#Base (the #tail v2)})
                 (node#= {#Hierarchy (the #root v1)}
                         {#Hierarchy (the #root v2)})))))))

(def node_mix
  (Mix Node)
  (implementation
   (def (mix $ init xs)
     (when xs
       {#Base base}
       (array.mix (function (_ _ item output) ($ item output))
                  init
                  base)
       
       {#Hierarchy hierarchy}
       (array.mix (function (_ _ node init') (mix $ init' node))
                  init
                  hierarchy)))))

(def .public mix
  (Mix Sequence)
  (implementation
   (def (mix $ init xs)
     (let [(open "[0]") node_mix]
       (mix $
            (mix $
                 init
                 {#Hierarchy (the #root xs)})
            {#Base (the #tail xs)})))))

(def .public monoid
  (All (_ a) (Monoid (Sequence a)))
  (implementation
   (def identity ..empty)
   
   (def (composite xs ys)
     (list#mix suffix xs (..list ys)))))

(def node_functor
  (Functor Node)
  (implementation
   (def (each $ xs)
     (when xs
       {#Base base}
       {#Base (array.each $ base)}
       
       {#Hierarchy hierarchy}
       {#Hierarchy (array.each (each $) hierarchy)}))))

(def .public functor
  (Functor Sequence)
  (implementation
   (def (each $ xs)
     [#level (the #level xs)
      #size (the #size xs)
      #root (let [ ... TODO: This binding was established to get around a compilation error. Fix and inline!
                  $ (at node_functor each $)]
              (|> xs (the #root) (array.each $)))
      #tail (|> xs (the #tail) (array.each $))])))

(def .public apply
  (Apply Sequence)
  (implementation
   (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)))))

(def .public monad
  (Monad Sequence)
  (implementation
   (def functor ..functor)

   (def in
     (|>> sequence))

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

(def .public reversed
  (All (_ a) (-> (Sequence a) (Sequence a)))
  (|>> ..list
       list.reversed
       (list#mix suffix ..empty)))

(with_template [<name> <array> <init> <op>]
  [(def .public <name>
     (All (_ a)
       (-> (Predicate a) (Sequence a) Bit))
     (let [help (is (All (_ a)
                      (-> (Predicate a) (Node a) Bit))
                    (function (help predicate node)
                      (when node
                        {#Base base}
                        (<array> predicate base)

                        {#Hierarchy hierarchy}
                        (<array> (help predicate) hierarchy))))]
       (function (<name> predicate sequence)
         (let [(open "_[0]") sequence]
           (<op> (help predicate {#Hierarchy _#root})
                 (help predicate {#Base _#tail}))))))]

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

(def .public (only when items)
  (All (_ a) (-> (-> a Bit) (Sequence a) (Sequence a)))
  (..mix (function (_ item output)
           (if (when item)
             (..suffix item output)
             output))
         ..empty
         items))

(def (one|node check items)
  (All (_ a b)
    (-> (-> a (Maybe b)) (Node a) (Maybe b)))
  (when items
    {#Base items}
    (array.one check items)
    
    {#Hierarchy items}
    (array.one (one|node check) items)))

(def .public (one check items)
  (All (_ a b)
    (-> (-> a (Maybe b)) (Sequence a) (Maybe b)))
  (when (let [... TODO: This binding was established to get around a compilation error. Fix and inline!
              check (..one|node check)]
          (|> items
              (the #root)
              (array.one check)))
    {.#None}
    (|> items
        (the #tail)
        (array.one check))
    
    output
    output))