Arrays with constant time indexing and slicing

This commit is contained in:
Nathan McCarty 2025-01-23 04:18:32 -05:00
parent 49525e43a1
commit 6d130cdc3b
3 changed files with 544 additions and 0 deletions

View file

@ -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

View file

@ -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

537
src/Array.md Normal file
View file

@ -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
```