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

(.module:
  [lux #*
   [abstract
    [monoid (#+ Monoid)]
    [functor (#+ Functor)]
    [equivalence (#+ Equivalence)]
    fold
    [predicate (#+ Predicate)]]
   [data
    ["." product]
    ["." maybe]
    [collection
     ["." list ("#@." fold)]]]
   [tool
    [compiler
     ["." host]]]])

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

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

(with-expansions [<elem-type> (primitive "java.lang.Object")
                  <array-type> (type (Array <elem-type>))]
  (def: #export (new size)
    (All [a] (-> Nat (Array a)))
    (`` (for {(~~ (static host.old))
              (:assume ("jvm anewarray" "(java.lang.Object )" size))

              (~~ (static host.jvm))
              (:assume
               (: (Array (primitive "java.lang.Object"))
                  ("jvm array new" size)))})))

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

              (~~ (static host.jvm))
              ("jvm array length" (:coerce <array-type> array))})))

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

                (~~ (static host.jvm))
                (let [value ("jvm array read" index (:coerce <array-type> array))]
                  (if ("jvm object null?" value)
                    #.None
                    (#.Some (:assume value))))}))
      #.None))

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

              (~~ (static host.jvm))
              (:assume
               ("jvm array write" index (:coerce <elem-type> value)
                (:coerce <array-type> array)))})))

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

                (~~ (static host.jvm))
                (write index (:assume (: <elem-type> ("jvm object null"))) 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 (occupied 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 (vacant array)
  {#.doc "Finds out how many cells in an array are vacant."}
  (All [a] (-> (Array a) Nat))
  (n/- (occupied array) (size array)))

(def: #export (filter! p xs)
  (All [a]
    (-> (-> a Bit) (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]
    (-> (-> a Bit) (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) (Array a) Bit))
     (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]
  )