{-# LANGUAGE GADTs #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE DataKinds #-}

module Agda.Utils.IndexedList where

import Data.Kind ( Type )
import Agda.Utils.Lens

-- | Existential wrapper for indexed types.
data Some :: (k -> Type) -> Type where
  Some :: f i -> Some f

-- | Unpacking a wrapped value.
withSome :: Some b -> (forall i. b i -> a) -> a
withSome :: forall {k} (b :: k -> *) a.
Some b -> (forall (i :: k). b i -> a) -> a
withSome (Some b i
x) forall (i :: k). b i -> a
f = forall (i :: k). b i -> a
f b i
x

-- | Lists indexed by a type-level list. A value of type @All p [x₁..xₙ]@ is a
--   sequence of values of types @p x₁@, .., @p xₙ@.
data All :: (x -> Type) -> [x] -> Type where
  Nil  :: All p '[]
  Cons :: p x -> All p xs -> All p (x ': xs)

-- | Constructing an indexed list from a plain list.
makeAll :: (a -> Some b) -> [a] -> Some (All b)
makeAll :: forall {x} a (b :: x -> *). (a -> Some b) -> [a] -> Some (All b)
makeAll a -> Some b
f [] = forall {k} (f :: k -> *) (xs :: k). f xs -> Some f
Some forall {x} (p :: x -> *). All p '[]
Nil
makeAll a -> Some b
f (a
x : [a]
xs) =
  case (a -> Some b
f a
x, forall {x} a (b :: x -> *). (a -> Some b) -> [a] -> Some (All b)
makeAll a -> Some b
f [a]
xs) of
    (Some b i
y, Some All b i
ys) -> forall {k} (f :: k -> *) (xs :: k). f xs -> Some f
Some (forall {x} (p :: x -> *) (xs :: x) (y :: [x]).
p xs -> All p y -> All p (xs : y)
Cons b i
y All b i
ys)

-- | Turning an indexed list back into a plain list.
forgetAll :: (forall x. b x -> a) -> All b xs -> [a]
forgetAll :: forall {x} (b :: x -> *) a (xs :: [x]).
(forall (x :: x). b x -> a) -> All b xs -> [a]
forgetAll forall (x :: x). b x -> a
f All b xs
Nil         = []
forgetAll forall (x :: x). b x -> a
f (Cons b x
x All b xs
xs) = forall (x :: x). b x -> a
f b x
x forall a. a -> [a] -> [a]
: forall {x} (b :: x -> *) a (xs :: [x]).
(forall (x :: x). b x -> a) -> All b xs -> [a]
forgetAll forall (x :: x). b x -> a
f All b xs
xs

-- | An index into a type-level list.
data Index :: [x] -> x -> Type where
  Zero :: Index (x ': xs) x
  Suc  :: Index xs x -> Index (y ': xs) x

-- | Indices are just natural numbers.
forgetIndex :: Index xs x -> Int
forgetIndex :: forall {x} (xs :: [x]) (x :: x). Index xs x -> Int
forgetIndex Index xs x
Zero    = Int
0
forgetIndex (Suc Index xs x
i) = Int
1 forall a. Num a => a -> a -> a
+ forall {x} (xs :: [x]) (x :: x). Index xs x -> Int
forgetIndex Index xs x
i

-- | Mapping over an indexed list.
mapWithIndex :: (forall x. Index xs x -> p x -> q x) -> All p xs -> All q xs
mapWithIndex :: forall {x} (xs :: [x]) (p :: x -> *) (q :: x -> *).
(forall (x :: x). Index xs x -> p x -> q x) -> All p xs -> All q xs
mapWithIndex forall (x :: x). Index xs x -> p x -> q x
f All p xs
Nil = forall {x} (p :: x -> *). All p '[]
Nil
mapWithIndex forall (x :: x). Index xs x -> p x -> q x
f (Cons p x
p All p xs
ps) = forall {x} (p :: x -> *) (xs :: x) (y :: [x]).
p xs -> All p y -> All p (xs : y)
Cons (forall (x :: x). Index xs x -> p x -> q x
f forall {x} (x :: x) (xs :: [x]). Index (x : xs) x
Zero p x
p) forall a b. (a -> b) -> a -> b
$ forall {x} (xs :: [x]) (p :: x -> *) (q :: x -> *).
(forall (x :: x). Index xs x -> p x -> q x) -> All p xs -> All q xs
mapWithIndex (forall (x :: x). Index xs x -> p x -> q x
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {x} (xs :: [x]) (x :: x) (y :: x).
Index xs x -> Index (y : xs) x
Suc) All p xs
ps

-- | If you have an index you can get a lens for the given element.
lIndex :: Index xs x -> Lens' (p x) (All p xs)
lIndex :: forall {x} (xs :: [x]) (x :: x) (p :: x -> *).
Index xs x -> Lens' (p x) (All p xs)
lIndex Index xs x
Zero    p x -> f (p x)
f (Cons p x
x All p xs
xs) = p x -> f (p x)
f p x
x           forall (m :: * -> *) a b. Functor m => m a -> (a -> b) -> m b
<&> \ p x
x  -> forall {x} (p :: x -> *) (xs :: x) (y :: [x]).
p xs -> All p y -> All p (xs : y)
Cons p x
x All p xs
xs
lIndex (Suc Index xs x
i) p x -> f (p x)
f (Cons p x
x All p xs
xs) = forall {x} (xs :: [x]) (x :: x) (p :: x -> *).
Index xs x -> Lens' (p x) (All p xs)
lIndex Index xs x
i p x -> f (p x)
f All p xs
xs forall (m :: * -> *) a b. Functor m => m a -> (a -> b) -> m b
<&> \ All p xs
xs -> forall {x} (p :: x -> *) (xs :: x) (y :: [x]).
p xs -> All p y -> All p (xs : y)
Cons p x
x All p xs
xs

-- | Looking up an element in an indexed list.
lookupIndex :: All p xs -> Index xs x -> p x
lookupIndex :: forall {x} (p :: x -> *) (xs :: [x]) (x :: x).
All p xs -> Index xs x -> p x
lookupIndex = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall {x} (xs :: [x]) (x :: x) (p :: x -> *).
Index xs x -> All p xs -> p x
ix
  where
    -- -Wincomplete-patterns fails for the other argument order!
    ix :: Index xs x -> All p xs -> p x
    ix :: forall {x} (xs :: [x]) (x :: x) (p :: x -> *).
Index xs x -> All p xs -> p x
ix Index xs x
Zero    (Cons p x
x All p xs
xs) = p x
x
    ix (Suc Index xs x
i) (Cons p x
x All p xs
xs) = forall {x} (xs :: [x]) (x :: x) (p :: x -> *).
Index xs x -> All p xs -> p x
ix Index xs x
i All p xs
xs

-- | All indices into an indexed list.
allIndices :: All p xs -> All (Index xs) xs
allIndices :: forall {x} (p :: x -> *) (xs :: [x]). All p xs -> All (Index xs) xs
allIndices = forall {x} (xs :: [x]) (p :: x -> *) (q :: x -> *).
(forall (x :: x). Index xs x -> p x -> q x) -> All p xs -> All q xs
mapWithIndex forall a b. a -> b -> a
const