(.module:
  [lux #*
   [control
    [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)))

(def: #export (new size)
  (All [a] (-> Nat (Array a)))
  (`` (for {(~~ (static host.jvm))
            (:assume ("jvm anewarray" "(java.lang.Object )" size))})))

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

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

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

    _
    #0))

(def: #export (write i x xs)
  (All [a]
    (-> Nat a (Array a) (Array a)))
  (`` (for {(~~ (static host.jvm))
            ("jvm aastore" xs i x)})))

(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 (delete i xs)
  (All [a]
    (-> Nat (Array a) (Array a)))
  (if (n/< (size xs) i)
    (`` (for {(~~ (static host.jvm))
              (write i (:assume ("jvm object null")) xs)}))
    xs))

(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]
                               #1

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

                               _
                               #0)))
                      #1
                      (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? #1 and]
  [any?   #0 or]
  )