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

(.module:
  [lux #*
   ["@" target]
   [abstract
    [monoid (#+ Monoid)]
    [functor (#+ Functor)]
    [equivalence (#+ Equivalence)]
    fold
    [predicate (#+ Predicate)]]
   [data
    ["." product]
    ["." maybe]
    [number
     ["n" nat]]
    [collection
     ["." list ("#\." fold)]]]])

(def: #export type-name "#Array")

(type: #export (Array a)
  {#.doc "Mutable arrays."}
  (#.Primitive ..type-name (#.Cons a #.Nil)))

(with-expansions [<index-type> (primitive "java.lang.Long")
                  <elem-type> (primitive "java.lang.Object")
                  <array-type> (type (Array <elem-type>))]
  (for {@.jvm
        (template: (!int value)
          (|> value
              (:coerce <index-type>)
              "jvm object cast"
              "jvm conversion long-to-int"))}
       (as-is))
  
  (def: #export (new size)
    (All [a] (-> Nat (Array a)))
    (for {@.old
          (:assume ("jvm anewarray" "(java.lang.Object )" size))

          @.jvm
          (|> size
              !int
              "jvm array new object"
              (: <array-type>)
              :assume)

          @.js
          ("js array new" size)}))

  (def: #export (size array)
    (All [a] (-> (Array a) Nat))
    (for {@.old
          ("jvm arraylength" array)

          @.jvm
          (|> array
              (:coerce <array-type>)
              "jvm array length object"
              "jvm conversion int-to-long"
              "jvm object cast"
              (: <index-type>)
              (:coerce Nat))

          @.js
          ("js array length" array)}))

  (def: #export (read index array)
    (All [a]
      (-> Nat (Array a) (Maybe a)))
    (if (n.< (size array) index)
      (for {@.old
            (let [value ("jvm aaload" array index)]
              (if ("jvm object null?" value)
                #.None
                (#.Some value)))

            @.jvm
            (let [value (|> array
                            (:coerce <array-type>)
                            ("jvm array read object" (!int index)))]
              (if ("jvm object null?" value)
                #.None
                (#.Some (:assume value))))

            @.js
            (let [output ("js array read" index array)]
              (if ("js object undefined?" output)
                #.None
                (#.Some output)))})
      #.None))

  (def: #export (write! index value array)
    (All [a]
      (-> Nat a (Array a) (Array a)))
    (for {@.old
          ("jvm aastore" array index value)

          @.jvm
          (|> array
              (:coerce <array-type>)
              ("jvm array write object" (!int index) (:coerce <elem-type> value))
              :assume)

          @.js
          ("js array write" index value array)}))

  (def: #export (delete! index array)
    (All [a]
      (-> Nat (Array a) (Array a)))
    (if (n.< (size array) index)
      (for {@.old
            (write! index (:assume ("jvm object null")) array)

            @.jvm
            (write! index (:assume (: <elem-type> ("jvm object null"))) array)

            @.js
            ("js array delete" index array)})
      array))
  )

(def: #export (contains? index array)
  (All [a]
    (-> Nat (Array a) Bit))
  (case (..read index array)
    (#.Some _)
    true

    _
    false))

(def: #export (update! index transform array)
  (All [a]
    (-> Nat (-> a a) (Array a) (Array a)))
  (case (read index array)
    #.None
    array

    (#.Some value)
    (write! index (transform value) array)))

(def: #export (upsert! index default transform array)
  (All [a]
    (-> Nat a (-> a a) (Array a) (Array a)))
  (write! index
          (|> array (read index) (maybe.default default) transform)
          array))

(def: #export (copy! length src-start src-array dest-start dest-array)
  (All [a]
    (-> Nat Nat (Array a) Nat (Array a)
        (Array a)))
  (if (n.= 0 length)
    dest-array
    (list\fold (function (_ offset target)
                 (case (read (n.+ offset src-start) src-array)
                   #.None
                   target
                   
                   (#.Some value)
                   (write! (n.+ offset dest-start) value target)))
               dest-array
               (list.indices length))))

(def: #export (occupancy array)
  {#.doc "Finds out how many cells in an array are occupied."}
  (All [a] (-> (Array a) Nat))
  (list\fold (function (_ idx count)
               (case (read idx array)
                 #.None
                 count
                 
                 (#.Some _)
                 (inc count)))
             0
             (list.indices (size array))))

(def: #export (vacancy array)
  {#.doc "Finds out how many cells in an array are vacant."}
  (All [a] (-> (Array a) Nat))
  (n.- (..occupancy array) (..size array)))

(def: #export (filter! p xs)
  (All [a]
    (-> (Predicate a) (Array a) (Array a)))
  (list\fold (function (_ idx xs')
               (case (read idx xs)
                 #.None
                 xs'

                 (#.Some x)
                 (if (p x)
                   xs'
                   (delete! idx xs'))))
             xs
             (list.indices (size xs))))

(def: #export (find p xs)
  (All [a]
    (-> (Predicate a) (Array a) (Maybe a)))
  (let [arr-size (size xs)]
    (loop [idx 0]
      (if (n.< arr-size idx)
        (case (read idx xs)
          #.None
          (recur (inc idx))
          
          (#.Some x)
          (if (p x)
            (#.Some x)
            (recur (inc idx))))
        #.None))))

(def: #export (find+ p xs)
  {#.doc "Just like 'find', but with access to the index of each value."}
  (All [a]
    (-> (-> Nat a Bit) (Array a) (Maybe [Nat a])))
  (let [arr-size (size xs)]
    (loop [idx 0]
      (if (n.< arr-size idx)
        (case (read idx xs)
          #.None
          (recur (inc idx))
          
          (#.Some x)
          (if (p idx x)
            (#.Some [idx x])
            (recur (inc idx))))
        #.None))))

(def: #export (clone xs)
  (All [a] (-> (Array a) (Array a)))
  (let [arr-size (size xs)]
    (list\fold (function (_ idx ys)
                 (case (read idx xs)
                   #.None
                   ys

                   (#.Some x)
                   (write! idx x ys)))
               (new arr-size)
               (list.indices arr-size))))

(def: #export (from-list xs)
  (All [a] (-> (List a) (Array a)))
  (product.right (list\fold (function (_ x [idx arr])
                              [(inc idx) (write! idx x arr)])
                            [0 (new (list.size xs))]
                            xs)))

(def: underflow Nat (dec 0))

(def: #export (to-list array)
  (All [a] (-> (Array a) (List a)))
  (loop [idx (dec (size array))
         output #.Nil]
    (if (n.= ..underflow idx)
      output
      (recur (dec idx)
             (case (read idx array)
               (#.Some head)
               (#.Cons head output)

               #.None
               output)))))

(def: #export (to-list' default array)
  (All [a] (-> a (Array a) (List a)))
  (loop [idx (dec (size array))
         output #.Nil]
    (if (n.= ..underflow idx)
      output
      (recur (dec idx)
             (#.Cons (maybe.default default (read idx array))
                     output)))))

(structure: #export (equivalence (^open ",\."))
  (All [a] (-> (Equivalence a) (Equivalence (Array a))))
  
  (def: (= xs ys)
    (let [sxs (size xs)
          sxy (size ys)]
      (and (n.= sxy sxs)
           (list\fold (function (_ idx prev)
                        (and prev
                             (case [(read idx xs) (read idx ys)]
                               [#.None #.None]
                               true

                               [(#.Some x) (#.Some y)]
                               (,\= x y)

                               _
                               false)))
                      true
                      (list.indices sxs))))))

(structure: #export monoid
  (All [a] (Monoid (Array a)))
  
  (def: identity (new 0))

  (def: (compose xs ys)
    (let [sxs (size xs)
          sxy (size ys)]
      (|> (new (n.+ sxy sxs))
          (copy! sxs 0 xs 0)
          (copy! sxy 0 ys sxs)))))

(structure: #export functor
  (Functor Array)
  
  (def: (map f ma)
    (let [arr-size (size ma)]
      (if (n.= 0 arr-size)
        (new arr-size)
        (list\fold (function (_ idx mb)
                     (case (read idx ma)
                       #.None
                       mb

                       (#.Some x)
                       (write! idx (f x) mb)))
                   (new arr-size)
                   (list.indices arr-size))
        ))))

(structure: #export fold
  (Fold Array)
  
  (def: (fold f init xs)
    (let [arr-size (size xs)]
      (loop [so-far init
             idx 0]
        (if (n.< arr-size idx)
          (case (read idx xs)
            #.None
            (recur so-far (inc idx))

            (#.Some value)
            (recur (f value so-far) (inc idx)))
          so-far)))))

(template [<name> <init> <op>]
  [(def: #export (<name> predicate array)
     (All [a]
       (-> (Predicate a) (Predicate (Array a))))
     (let [size (..size array)]
       (loop [idx 0]
         (if (n.< size idx)
           (case (..read idx array)
             (#.Some value)
             (<op> (predicate value)
                   (recur (inc idx)))
             
             #.None
             (recur (inc idx)))
           <init>))))]

  [every? true  and]
  [any?   false or]
  )