path: root/stdlib/source/lux/data/struct/vector.lux

##  Copyright (c) Eduardo Julian. All rights reserved.
##  This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0.
##  If a copy of the MPL was not distributed with this file,
##  You can obtain one at http://mozilla.org/MPL/2.0/.

(;module:
  lux
  (lux (control functor
                applicative
                monad
                eq
                monoid
                fold)
       (data maybe
             (struct [list "List/" Fold<List> Functor<List> Monoid<List>]
                     [array #+ Array "Array/" Functor<Array> Fold<Array>])
             [bit]
             [number "Int/" Number<Int>]
             [product])
       [compiler #+ with-gensyms]
       (macro [ast]
              ["s" syntax #+ syntax: Syntax])
       [pipe]
       ))

## This implementation of vectors is based on Clojure's
## PersistentVector implementation.

## [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;<< branching-exponent +1))

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

(def: branch-idx
  (-> Index Index)
  (bit;& branch-idx-mask))

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

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

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

(def: (new-tail singleton)
  (All [a] (-> a (Base a)))
  (|> (: (Base ($ +0))
         (array;new +1))
      (array;put +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;>>> level (n.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;get 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;put sub-idx sub-node))))

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

(def: (put' level idx val hierarchy)
  (All [a] (-> Level Index a (Hierarchy a) (Hierarchy a)))
  (let [sub-idx (branch-idx (bit;>>> level idx))]
    (case (array;get sub-idx hierarchy)
      (#;Some (#Hierarchy sub-node))
      (|> (array;clone hierarchy)
          (array;put sub-idx (#Hierarchy (put' (level-down level) idx val sub-node))))

      (^=> (#;Some (#Base base))
           (n.= +0 (level-down level)))
      (|> (array;clone hierarchy)
          (array;put sub-idx (|> (array;clone base)
                                 (array;put (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;>>> level (n.- +2 size)))]
    (cond (n.= +0 sub-idx)
          #;None

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

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

          ## Else...
          (|> (array;clone hierarchy)
              (array;remove 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 (lambda [sub acc] (List/append (to-list' sub) acc))
                   #;Nil))))

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

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

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

(def: #export (add val vec)
  (All [a] (-> a (Vector a) (Vector 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 n.inc)
          (update@ #tail (expand-tail val)))
      ## Otherwise, push tail into the tree
      ## --------------------------------------------------------
      ## Will the root experience an overflow with this addition?
      (|> (if (n.> (bit;<< (get@ #level vec) +1)
                   (bit;>>> 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 []))
                                (array;put +0 (#Hierarchy (get@ #root vec)))
                                (array;put +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 Vector and grow a new
          ## tail with the new element as it's sole member.
          (update@ #size n.inc)
          (set@ #tail (new-tail val)))
      )))

(def: (base-for idx vec)
  (All [a] (-> Index (Vector 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;get (branch-idx (bit;>>> level idx)) hierarchy)]
            [true (#;Some (#Hierarchy sub))]
            (recur (level-down level) sub)

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

            [_ #;None]
            #;None

            _
            (error! "Incorrect vector structure."))))
      #;None)))

(def: #export (at idx vec)
  (All [a] (-> Nat (Vector a) (Maybe a)))
  (do Monad<Maybe>
    [base (base-for idx vec)]
    (array;get (branch-idx idx) base)))

(def: #export (put idx val vec)
  (All [a] (-> Nat a (Vector a) (Vector 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 (: (-> (Base ($ +0)) (Base ($ +0)))
                              (|>. array;clone (array;put (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) (Vector a) (Vector a)))
  (case (at idx vec)
    (#;Some val)
    (put idx (f val) vec)

    #;None
    vec))

(def: #export (pop vec)
  (All [a] (-> (Vector a) (Vector 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 (n.dec (array;size old-tail))]
        (|> vec
            (update@ #size n.dec)
            (set@ #tail (|> (array;new new-tail-size)
                            (array;copy new-tail-size +0 old-tail +0)))))
      (default (undefined)
        (do Monad<Maybe>
          [new-tail (base-for (n.- +2 vec-size) vec)
           #let [[level' root'] (: [Level (Hierarchy ($ +0))]
                                   (let [init-level (get@ #level vec)]
                                     (loop [level init-level
                                            root (: (Hierarchy ($ +0))
                                                    (default (new-hierarchy [])
                                                      (pop-tail vec-size init-level (get@ #root vec))))]
                                       (if (n.> branching-exponent level)
                                         (case [(array;get +1 root) (array;get +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 n.dec)
                    (set@ #level level')
                    (set@ #root root')
                    (set@ #tail new-tail))))))
    ))

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

(def: #export (from-list list)
  (All [a] (-> (List a) (Vector a)))
  (List/fold add
             (: (Vector ($ +0))
                empty)
             list))

(def: #export (member? a/Eq vec val)
  (All [a] (-> (Eq a) (Vector a) a Bool))
  (list;member? a/Eq (to-list vec) val))

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

## [Syntax]
(syntax: #export (vector [elems (s;some s;any)])
  {#;doc (doc "Vector literals."
              (vector 10 20 30 40))}
  (wrap (list (` (from-list (list (~@ elems)))))))

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

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

(struct: _ (Fold Node)
  (def: (fold f init xs)
    (case xs
      (#Base base)
      (Array/fold f init base)
      
      (#Hierarchy hierarchy)
      (Array/fold (lambda [node init'] (fold f init' node))
                  init
                  hierarchy))
    ))

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

(struct: #export Monoid<Vector> (All [a]
                                  (Monoid (Vector a)))
  (def: unit empty)
  (def: (append xs ys)
    (List/fold add xs (to-list ys))))

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

(struct: #export _ (Functor Vector)
  (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))
     }))

(struct: #export _ (Applicative Vector)
  (def: functor Functor<Vector>)

  (def: (wrap x)
    (vector x))
  
  (def: (apply ff fa)
    (let [(^open) Functor<Vector>
          (^open) Fold<Vector>
          (^open) Monoid<Vector>
          results (map (lambda [f] (map f fa))
                       ff)]
      (fold append unit results)))
  )

(struct: #export _ (Monad Vector)
  (def: applicative Applicative<Vector>)

  (def: join
    (let [(^open) Fold<Vector>
          (^open) Monoid<Vector>]
      (fold (lambda [post pre] (append pre post)) unit)))
  )