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 Provide views that enable recursively pattern matching numbers as lists of
digits, in both ascending and descending order of significance. 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 ## Index of years and days
- 2015 - 2015

View file

@ -16,6 +16,7 @@ authors = "Nathan McCarty"
depends = base depends = base
, contrib , contrib
, structures , structures
, tailrec
, eff , eff
, elab-util , elab-util
, ansi , ansi
@ -28,6 +29,7 @@ modules = Runner
, Util , Util
, Util.Eff , Util.Eff
, Util.Digits , Util.Digits
, Array
-- main file (i.e. file to load at REPL) -- main file (i.e. file to load at REPL)
main = Main 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
```