Array Wrappers

module Array

import Data.IOArray
import Data.Vect
import Data.List.Lazy
import Data.Fin
import Data.Iterable
import Decidable.Equality
import Control.WellFounded

%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.

%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.

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.

allFins : (n : Nat) -> LazyList (Fin n)
allFins 0 = []
allFins (S k) = FZ :: map FS (allFins k)

LazyLists have a size

Sized (LazyList e) where
  size [] = 0
  size (x :: xs) = S (size xs)

Convert a Vect to a LazyList

vectLazy : Vect n e -> LazyList e
vectLazy [] = []
vectLazy (x :: xs) = x :: vectLazy xs

LazyLists are iterable

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

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.

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

  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

  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

  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.

  %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

  -- 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 IArrays as if they were lists

  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

  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

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

  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.

  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

  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.

  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.

  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.

  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

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