From 6d130cdc3b4bf6c96ed76991d52c7c7dc2e5c900 Mon Sep 17 00:00:00 2001
From: Nathan McCarty <thatonelutenist@stranger.systems>
Date: Thu, 23 Jan 2025 04:18:32 -0500
Subject: [PATCH] Arrays with constant time indexing and slicing

---
 README.md    |   5 +
 advent.ipkg  |   2 +
 src/Array.md | 537 +++++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 544 insertions(+)
 create mode 100644 src/Array.md

diff --git a/README.md b/README.md
index e2cf08f..d8bd7e3 100644
--- a/README.md
+++ b/README.md
@@ -51,6 +51,11 @@ solution.
     Provide views that enable recursively pattern matching numbers as lists of
     digits, in both ascending and descending order of significance.
 
+- [Array](src/Array.md)
+
+  Provider wrappers over the standard library `IOArray` type to make them more
+  ergonomic to use.
+
 ## Index of years and days
 
 - 2015
diff --git a/advent.ipkg b/advent.ipkg
index 189172e..23c2130 100644
--- a/advent.ipkg
+++ b/advent.ipkg
@@ -16,6 +16,7 @@ authors = "Nathan McCarty"
 depends = base
         , contrib
         , structures
+        , tailrec
         , eff
         , elab-util
         , ansi
@@ -28,6 +29,7 @@ modules = Runner
         , Util
         , Util.Eff
         , Util.Digits
+        , Array
 
 -- main file (i.e. file to load at REPL)
 main = Main
diff --git a/src/Array.md b/src/Array.md
new file mode 100644
index 0000000..98bf078
--- /dev/null
+++ b/src/Array.md
@@ -0,0 +1,537 @@
+# Array Wrappers
+
+```idris
+module Array
+
+import Data.IOArray
+import Data.Vect
+import Data.List.Lazy
+import Data.Fin
+import Data.Iterable
+import Decidable.Equality
+import Control.WellFounded
+```
+
+<!-- idris
+import System
+-->
+
+```idris
+%default total
+%hide Prelude.Types.elem
+```
+
+Provide some wrappers over the standard library IO arrays, using plenty of
+unsafe operations along the way.
+
+## Internal Utility Functions
+
+Unsafely unwrap a `Maybe`, making the assumption that it is a `Just`.
+
+```idris
+%unsafe
+unwrapUnsafe : Maybe a -> a
+unwrapUnsafe x = assert_total $ unwrapUnsafe' x
+  where 
+    partial
+    unwrapUnsafe' : Maybe a -> a
+    unwrapUnsafe' (Just x) = x
+```
+
+Insert an element into an `IOArray` at the index described by the given `Fin`.
+
+```idris
+insertPair : IOArray e -> (e, Fin n) -> IO ()
+insertPair arr (elem, idx) = do
+  let idx : Int = cast . finToInteger $ idx
+  _ <- writeArray arr idx elem
+  pure ()
+```
+
+### Tail safe monadic recursion with lazy lists
+
+Generate all of the fins of a given size, lazily.
+
+```idris
+allFins : (n : Nat) -> LazyList (Fin n)
+allFins 0 = []
+allFins (S k) = FZ :: map FS (allFins k)
+```
+
+LazyLists have a size
+
+```idris
+Sized (LazyList e) where
+  size [] = 0
+  size (x :: xs) = S (size xs)
+```
+
+Convert a `Vect` to a `LazyList`
+
+```idris
+vectLazy : Vect n e -> LazyList e
+vectLazy [] = []
+vectLazy (x :: xs) = x :: vectLazy xs
+```
+
+LazyLists are iterable
+
+```idris
+refl : {k : Nat} -> LTE k k
+refl = reflexive {x = k}
+
+Iterable (LazyList a) a where
+  iterM accum done ini seed =
+    trSized seed ini $
+      \case Nil    => pure . Done . done
+            h :: t => map (Cont t refl) . accum h
+```
+
+## Immutable arrays
+
+<!-- idris
+namespace Immutable
+-->
+
+Immutable arrays created from vectors that provide constant time indexing and
+slicing operations.
+
+As an invariant of the the type, every "slot" in the inner `IOArray` must always
+be a `Just`, unsafe operations are performed that depend on that.
+
+```idris
+  ||| An immutable array of length `n`
+  export
+  data IArray : Nat -> Type -> Type where
+    MkIArray : (size : Nat) -> (offset : Int) -> (inner : IOArray e) 
+      -> IArray size e
+  %name IArray xs, ys, zs
+```
+
+### Basic Interface
+
+Indexing operation
+
+```idris
+  export
+  index : (idx : Fin n) -> IArray n e -> e
+  index idx (MkIArray size offset inner) = unsafePerformIO index' 
+    where
+      index' : IO e
+      index' = do
+        -- Convert the index to an Int for use in the IO array
+        let idx : Int = cast . finToInteger $ idx
+        map unwrapUnsafe $ readArray inner (idx + offset) 
+```
+
+Getting the length of an `IArray`, as well as a proof of correctness for the
+length function
+
+```idris
+  export
+  length : IArray n e -> Nat
+  length (MkIArray n _ inner) = n
+  
+  export
+  lengthCorrect : (array : IArray n e) -> length array = n
+  lengthCorrect (MkIArray n _ inner) = Refl
+```
+
+Basic conversion operations
+
+```idris
+  export
+  fromVect : {n : Nat} -> Vect n e -> IArray n e
+  fromVect xs = unsafePerformIO fromVect'
+    where
+      fromVect' : IO $ IArray n e
+      fromVect' = do
+        let pairs = zip (vectLazy xs) (allFins n)
+        array <- newArray (cast n)
+        forM_ (insertPair array) pairs
+        pure $ MkIArray n 0 array
+  
+  export
+  replicate : (n : Nat) -> e -> IArray n e
+  replicate n x = unsafePerformIO replicate'
+    where
+      insertR : IOArray e -> Fin n -> IO ()
+      insertR arr idx = do
+        let idx : Int = cast . finToInteger $ idx
+        _ <- writeArray arr idx x
+        pure ()
+      replicate' : IO $ IArray n e
+      replicate' = do
+        array <- newArray (cast n)
+        forM_ (insertR array) (Array.allFins n)
+        pure $ MkIArray n 0 array
+
+  export
+  toVect : IArray n e -> Vect n e
+  toVect xs = 
+    let xs' : IArray (length xs) e = rewrite lengthCorrect xs in xs
+        values : Vect _ e = map (flip index $ xs') (allFins _)
+    in rewrite sym $ lengthCorrect xs in values
+  
+  arrayToList : IArray n e -> List e
+  arrayToList xs = 
+    let xs' : IArray (length xs) e = rewrite lengthCorrect xs in xs
+    in map (flip index $ xs') (allFins (length xs))
+```
+
+Unsafely create an array from a lazy list of values, for internal use only. Will
+cause crashes if the provided `LazyList` isn't the right length.
+
+```idris
+  %unsafe
+  unsafeFromLazy : {n : Nat} -> LazyList e -> IO (IArray n e)
+  unsafeFromLazy xs = do
+    array <- newArray (cast n)
+    let pairs = zip xs (allFins n)
+    forM_ (insertPair array) pairs
+    pure $ MkIArray _ 0 array
+```
+
+Typical drop and take operations, implemented in constant time
+
+```idris
+  -- TODO: Minimize whatever compiler bug is requiring us to use this
+  newOffset : Nat -> Int -> Int
+  newOffset k i = cast k + i
+
+  export
+  drop : (n : Nat) -> IArray (n + m) e -> IArray m e
+  drop n xs@(MkIArray _ offset inner) = 
+    believe_me $ MkIArray (length xs `minus` n) (newOffset n offset) inner
+
+  export
+  drop' : (n : Nat) -> IArray l e -> IArray (minus l n) e
+  drop' n (MkIArray l offset inner) = 
+    if n >= l
+      then believe_me $ fromVect [] {e}
+      else MkIArray (l `minus` n) (newOffset n offset) inner
+  
+  export
+  take : (n : Nat) -> IArray (n + m) e -> IArray n e
+  take n (MkIArray _ offset inner) = MkIArray n offset inner
+```
+
+A view for pattern matching on `IArray`s as if they were lists
+
+```idris
+  namespace View
+    public export
+    data AsList : IArray n e -> Type where
+      Nil : {0 tail : IArray 0 e} -> AsList tail
+      (::) : {0 xs : IArray (S n) e} -> {rest : IArray n e} 
+        -> (head : e) -> (tail : Lazy (AsList rest)) -> AsList xs
+    
+    public export 
+    asList : (xs : IArray n e) -> AsList xs
+    asList (MkIArray 0 offset inner) = []
+    asList xs@(MkIArray (S k) offset inner) = 
+      let head = 0 `index` xs
+          rest = drop 1 xs
+      in head :: (asList rest)
+```
+
+Convert to a `LazyList` using our view
+
+```idris
+  export
+  toLazy : IArray n e -> LazyList e
+  toLazy xs with (asList xs)
+    toLazy xs | [] = []
+    toLazy xs | (head :: tail) = 
+      head :: toLazy _ | tail
+```
+
+Typical filtering and finding interface
+
+```idris
+  -- Don't know if this one is really worth optimizing
+  export
+  filter : (e -> Bool) -> IArray n e -> (p : Nat ** IArray p e)
+  filter f xs = 
+    let (_ ** ps) = filter f (toVect xs)
+    in (_ ** fromVect ps)
+ 
+  export
+  find : (e -> Bool) -> IArray n e -> Maybe e
+  find f xs with (asList xs)
+    find f xs | [] = Nothing
+    find f xs | (head :: tail) = 
+      if f head
+        then Just head
+        else find f _ | tail
+ 
+  export
+  findIndex : (e -> Bool) -> IArray n e -> Maybe (Fin n)
+  findIndex f xs = 
+    let indexed = zip (toLazy xs) (allFins (length xs))
+    in findIndex' f indexed
+    where
+      findIndex' : (e -> Bool) -> LazyList (e, Fin (length xs)) -> Maybe (Fin n)
+      findIndex' f [] = Nothing
+      findIndex' f ((e, idx) :: ys) = 
+        if f e
+          then rewrite sym $ lengthCorrect xs in Just idx
+          else findIndex' f ys
+ 
+  export
+  findIndices : (e -> Bool) -> IArray n e -> List (Fin n)
+  findIndices f xs@(MkIArray n _ _) = 
+    let pairs = zip (toLazy xs) (allFins n)
+    in toList . map snd . filter (f . fst) $ pairs
+```
+
+### Interface Implementations
+
+Provide a `Foldable` implementation by performing indexing within the context of
+`unsafePerformIO`.
+
+```idris
+  export
+  Foldable (IArray n) where
+    foldl f acc_val (MkIArray n offset inner) = unsafePerformIO $ foldl'
+      where 
+        liftF : Fin n -> acc -> IO acc
+        liftF idx acc_val = do
+          let idx : Int = cast . finToInteger $ idx
+          val <- map unwrapUnsafe $ readArray inner (idx + offset)
+          pure $ f acc_val val
+        foldl' : IO acc
+        foldl' = iterM liftF id acc_val (Array.allFins n)
+    foldr f acc_val (MkIArray n offset inner) = unsafePerformIO $ foldr'
+      where
+        liftF : Fin n -> acc -> IO acc
+        liftF idx acc_val = do
+          let idx : Int = cast . finToInteger $ idx
+          val <- map unwrapUnsafe $ readArray inner (idx + offset)
+          pure $ f val acc_val
+        foldr' : IO acc
+        -- TODO: Make a lazy reverse allFins function so we aren't instantiating
+        -- the full index list here
+        foldr' = iterM liftF id acc_val (reverse $ List.allFins n)
+    toList = arrayToList
+```
+
+Provide implementations of `Eq`, `Ord`, and `DecEq` in terms of our `AsList`
+view. The `DecEq` implementation requires use of unsafe `believe_me`, as arrays
+are primitive types the compiler has no insight into the structure of.
+
+```idris
+  export
+  Eq e => Eq (IArray n e) where
+    xs == ys = (toLazy xs) == (toLazy ys)
+  
+  export
+  Ord e => Ord (IArray n e) where
+    compare xs ys = compare (toLazy xs) (toLazy ys)
+  
+  -- TODO : Optimize IArray DecEq
+  decEq' : DecEq e => {m : Nat} -> (xs, ys : IArray m e) -> Dec (xs = ys)
+  decEq' xs ys with (asList xs, asList ys)
+    decEq' xs ys | ([], []) = believe_me $ the (xs = xs) Refl
+    decEq' xs ys | ((h_x :: t_x), (h_y :: t_y)) = 
+      case decEq h_x h_y of
+        Yes prf => 
+          believe_me $ decEq' _ _ | (t_x, t_y)
+        No contra => believe_me contra
+
+  export
+  DecEq e => DecEq (IArray n e) where
+    decEq xs@(MkIArray n _ _) ys = decEq' xs ys
+```
+
+Provide `Functor`, `Applicative`, `Monad`, and `Traverseable` implementations by
+building a new array in an `IO` context
+
+```idris
+  map' : (f : a -> b) -> IArray n a -> IO (IArray n b)
+  map' f xs@(MkIArray n _ _) = do
+    array <- newArray (cast n)
+    let pairs = zip (map f . toLazy $ xs) (Array.allFins n)
+    forM_ (insertPair array) pairs
+    pure $ MkIArray n 0 array
+
+  export
+  Functor (IArray k) where
+    map f xs = unsafePerformIO $ map' f xs
+  
+  apply' : {n : Nat} -> IArray n (a -> b) -> IArray n a -> IO (IArray n b)
+  apply' fs xs = do
+    array <- newArray (cast n)
+    let applied = map (\(f, x) => f x) $ zip (toLazy fs) (toLazy xs)
+    let pairs = zip applied (allFins n)
+    forM_ (insertPair array) pairs
+    pure $ MkIArray n 0 array
+ 
+  -- Applicative requires the length of the array to be available at runtime at
+  -- the type level for `replicate`
+  export
+  {k : Nat} -> Applicative (IArray k) where
+    pure = replicate _
+    fs <*> xs = unsafePerformIO $ apply' fs xs
+
+  -- Like `Vect`'s `join`, this takes the diagonal elements
+  export
+  {k : Nat} -> Monad (IArray k) where 
+    join xs = 
+      let lazys = join' . map toLazy $ toLazy xs
+      in assert_total $ unsafePerformIO $ unsafeFromLazy lazys
+      where
+        covering
+        join' : LazyList (LazyList a) -> LazyList a
+        join' [] = []
+        join' ([] :: y) = []
+        join' ((x :: _) :: xs) = 
+          x :: join' (map (drop 1) xs)
+    -- join xs = fromVect . join . map toVect . toVect $ xs
+ 
+  -- TODO: Maybe take another pass at optimizing this? 
+  export
+  Traversable (IArray k) where 
+    traverse fun xs@(MkIArray k _ _) = 
+      map (unsafePerformIO . unsafeFromLazy . fromList) 
+      . Prelude.traverse fun 
+      . toList 
+      $ xs
+```
+
+Provide a `Show` implementation in terms of our `AsList` view through the
+`asLazy` wrapper.
+
+```idris
+  export
+  Show e => Show (IArray n e) where
+    show xs = 
+      let elements = map show $ toLazy xs
+      in case elements of
+        [] => "[]"
+        (x :: xs) => "[\{join' xs x}]"
+      where
+        join' : LazyList String -> (acc : String) -> String
+        join' [] acc = acc
+        join' (x :: []) acc = "\{acc}, \{x}"
+        join' (x :: (y :: xs)) acc = 
+          join' (y :: xs) "\{acc}, \{x}"
+```
+
+Provide a `Zippable` implementation by roundtripping through a `Vect`. This
+isn't the most efficient, but this provides a workable starter implementation.
+
+```idris
+  export
+  Zippable (IArray k) where 
+    zipWith f xs@(MkIArray k _ _) ys = 
+      unsafePerformIO . unsafeFromLazy $ zipWith f (toLazy xs) (toLazy ys)
+    zipWith3 f xs@(MkIArray k _ _) ys zs = 
+      unsafePerformIO . unsafeFromLazy $ zipWith3 f (toLazy xs) (toLazy ys) (toLazy zs)
+    unzipWith f xs@(MkIArray k _ _) = 
+      let (xs, ys) = unzipWith f (toLazy xs)
+      in (unsafePerformIO . unsafeFromLazy $ xs, 
+          unsafePerformIO . unsafeFromLazy $ ys)
+    unzipWith3 f xs@(MkIArray k _ _) = 
+      let (xs, ys, zs) = unzipWith3 f (toLazy xs)
+      in (unsafePerformIO . unsafeFromLazy $ xs, 
+          unsafePerformIO . unsafeFromLazy $ ys, 
+          unsafePerformIO . unsafeFromLazy $ zs)
+```
+
+Provide `Semigroup` and `Monoid` implementations, matching those for `Vect`, in
+terms of our other interfaces.
+
+```idris
+  export
+  Semigroup a => Semigroup (IArray k a) where 
+    xs <+> ys = zipWith (<+>) xs ys
+  
+  -- Monoid requires the length argument to be available at runtime for 
+  -- `replicate` to work
+  export
+  {k : Nat} -> Monoid a => Monoid (IArray k a) where 
+    neutral = replicate _ neutral
+```
+
+## Unit tests
+
+### IArray
+
+```idris
+-- @@test IArray RoundTrip
+iRoundTrip : IO Bool
+iRoundTrip = do
+  let test_vect : Vect _ Nat = [1, 2, 3, 4, 5]
+  putStrLn "test_vect: \{show test_vect}"
+  let xs = fromVect test_vect
+  putStrLn "xs: \{show xs}"
+  let round_tripped = toVect xs 
+  putStrLn "round_tripped: \{show round_tripped}"
+  pure $ test_vect == round_tripped
+
+-- @@test IArray drop/take
+iDropTake : IO Bool
+iDropTake = do
+  let test = fromVect [1, 2, 3, 4, 5]
+  putStrLn "test: \{show test}"
+  let dropped = drop 2 test
+  putStrLn "dropped: \{show dropped}"
+  let taked = take 3 test
+  putStrLn "taked: \{show taked}"
+  pure $ (toVect dropped) == [3, 4, 5] && (toVect taked) == [1, 2, 3]
+
+-- @@test IArray toLazy
+iLazy : IO Bool
+iLazy = do
+  let test_vect : Vect _ Nat = [1, 2, 3, 4, 5]
+  let test_array = fromVect test_vect
+  putStrLn "test_array: \{show test_array}"
+  let output_lazy = toLazy test_array
+  putStrLn "output_lazy: \{show output_lazy}"
+  pure $ vectLazy test_vect == output_lazy
+
+-- @@test IArray foldl/foldr
+iFold : IO Bool
+iFold = do
+  let test = fromVect [1, 2, 3, 4, 5]
+  putStrLn "test: \{show test}"
+  let foldl_test = foldl (flip (::)) [] test
+  putStrLn "foldl_test: \{show foldl_test}"
+  let foldr_test = foldr (::) [] test
+  putStrLn "foldr_test: \{show foldr_test}"
+  pure $ foldl_test == [5, 4, 3, 2, 1] && foldr_test == [1, 2, 3, 4, 5]
+
+-- @@test IArray functor
+iFunctor : IO Bool
+iFunctor = do
+  let test = fromVect [1, 2, 3, 4, 5]
+  putStrLn "test: \{show test}"
+  let output = map (+ 1) test
+  putStrLn "output: \{show output}"
+  pure $ toVect output == [2, 3, 4, 5, 6]
+  
+-- @@test IArray show
+iShow : IO Bool
+iShow = do
+  let test = fromVect [1, 2, 3, 4, 5]
+  putStrLn "test: \{show test}"
+  let empty : IArray _ Nat = fromVect []
+  putStrLn "empty: \{show empty}"
+  let singleton = fromVect [1]
+  putStrLn "singleton: \{show singleton}"
+  pure $ 
+    show test == "[1, 2, 3, 4, 5]"  
+    && show empty == "[]"  
+    && show singleton == "[1]"  
+
+-- @@test IArray monad
+iMonad : IO Bool
+iMonad = do
+  let test_vect : Vect 3 (Vect 3 Nat) = [[1, 2, 3], [4, 5, 6], [7, 8 ,9]]
+  let test_array = fromVect $ map fromVect test_vect
+  let vect_out = join test_vect
+  let array_out = join test_array
+  putStrLn "vect: \{show vect_out} array: \{show array_out}"
+  pure $ toVect array_out == vect_out
+```