{-
   Copyright 2016, Dominic Orchard, Andrew Rice, Mistral Contrastin, Matthew Danish

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
-}

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators #-}

module Camfort.Helpers.Vec
  (
    -- * Datatypes
    EqT(..)
  , ExistsEqT(..)
  , Nat(..)
  , Natural(..)
  , Vec(..)
  , VecBox(..)
  , VecList(..)
  -- ** Vector Operations
  , (!!)
  , findIndex
  , fromList
  , fromLists
  , length
  , lengthN
  , proveEqSize
  , proveNonEmpty
  , replace
  , toList
  , zip
  , zipWith
  ) where

import Prelude hiding (length, zip, zipWith, take, drop, (!!))

import Data.Proxy

import Unsafe.Coerce

data Nat = Z | S Nat

-- Indexed natural number type
data Natural (n :: Nat) where
     Zero :: Natural 'Z
     Succ :: Natural n -> Natural ('S n)

deriving instance Show (Natural n)

data NatBox where NatBox :: Natural n -> NatBox
deriving instance Show NatBox

class IsNatural (n :: Nat) where
   fromNat :: Proxy n -> Int

instance IsNatural 'Z where
   fromNat :: Proxy 'Z -> Int
fromNat Proxy 'Z
Proxy = Int
0
instance IsNatural n => IsNatural ('S n) where
   fromNat :: Proxy ('S n) -> Int
fromNat Proxy ('S n)
Proxy = Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Proxy n -> Int
forall (n :: Nat). IsNatural n => Proxy n -> Int
fromNat (Proxy n
forall k (t :: k). Proxy t
Proxy :: Proxy n)

-- Indexed vector type
data Vec (n :: Nat) a where
     Nil :: Vec 'Z a
     Cons :: a -> Vec n a -> Vec ('S n) a

length :: Vec n a -> Int
length :: Vec n a -> Int
length Vec n a
Nil = Int
0
length (Cons a
_ Vec n a
xs) = Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Vec n a -> Int
forall (n :: Nat) a. Vec n a -> Int
length Vec n a
xs

lengthN :: Vec n a -> Natural n
lengthN :: Vec n a -> Natural n
lengthN Vec n a
Nil = Natural n
Natural 'Z
Zero
lengthN (Cons a
_ Vec n a
xs) = Natural n -> Natural ('S n)
forall (n :: Nat). Natural n -> Natural ('S n)
Succ (Natural n -> Natural ('S n)) -> Natural n -> Natural ('S n)
forall a b. (a -> b) -> a -> b
$ Vec n a -> Natural n
forall (n :: Nat) a. Vec n a -> Natural n
lengthN Vec n a
xs

instance Functor (Vec n) where
  fmap :: (a -> b) -> Vec n a -> Vec n b
fmap a -> b
_ Vec n a
Nil         = Vec n b
forall a. Vec 'Z a
Nil
  fmap a -> b
f (Cons a
x Vec n a
xs) = b -> Vec n b -> Vec ('S n) b
forall a (n :: Nat). a -> Vec n a -> Vec ('S n) a
Cons (a -> b
f a
x) ((a -> b) -> Vec n a -> Vec n b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f Vec n a
xs)

deriving instance Eq a => Eq (Vec n a)

instance Ord a => Ord (Vec n a) where
    Vec n a
Nil         <= :: Vec n a -> Vec n a -> Bool
<= Vec n a
_ = Bool
True
    (Cons a
x Vec n a
xs) <= (Cons a
y Vec n a
ys) | Vec n a
xs Vec n a -> Vec n a -> Bool
forall a. Eq a => a -> a -> Bool
== Vec n a
Vec n a
ys = a
x a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
y
                               | Bool
otherwise = Vec n a
xs Vec n a -> Vec n a -> Bool
forall a. Ord a => a -> a -> Bool
<= Vec n a
Vec n a
ys
instance Show a => Show (Vec n a) where
  show :: Vec n a -> String
show Vec n a
xs = String
"<" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Vec n a -> String
forall (n' :: Nat) a'. Show a' => Vec n' a' -> String
showV Vec n a
xs String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
">"
    where
      showV :: forall n' a' . Show a' => Vec n' a' -> String
      showV :: Vec n' a' -> String
showV Vec n' a'
Nil          = String
""
      showV (Cons a'
x Vec n a'
Nil)  = a' -> String
forall a. Show a => a -> String
show a'
x
      showV (Cons a'
x Vec n a'
xs')  = a' -> String
forall a. Show a => a -> String
show a'
x String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"," String -> ShowS
forall a. [a] -> [a] -> [a]
++ Vec n a' -> String
forall (n' :: Nat) a'. Show a' => Vec n' a' -> String
showV Vec n a'
xs'

instance Foldable (Vec n) where
  foldr :: (a -> b -> b) -> b -> Vec n a -> b
foldr a -> b -> b
_ b
acc Vec n a
Nil = b
acc
  foldr a -> b -> b
f b
acc (Cons a
x Vec n a
xs) = (a -> b -> b) -> b -> Vec n a -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> b -> b
f (a -> b -> b
f a
x b
acc) Vec n a
xs

zipWith :: (a -> b -> c) -> Vec n a -> Vec n b -> Vec n c
zipWith :: (a -> b -> c) -> Vec n a -> Vec n b -> Vec n c
zipWith a -> b -> c
_ Vec n a
Nil Vec n b
Nil = Vec n c
forall a. Vec 'Z a
Nil
zipWith a -> b -> c
f (Cons a
x Vec n a
xs) (Cons b
y Vec n b
ys) = c -> Vec n c -> Vec ('S n) c
forall a (n :: Nat). a -> Vec n a -> Vec ('S n) a
Cons (a -> b -> c
f a
x b
y) ((a -> b -> c) -> Vec n a -> Vec n b -> Vec n c
forall a b c (n :: Nat).
(a -> b -> c) -> Vec n a -> Vec n b -> Vec n c
zipWith a -> b -> c
f Vec n a
xs Vec n b
Vec n b
ys)

zip :: Vec n a -> Vec n b -> Vec n (a,b)
zip :: Vec n a -> Vec n b -> Vec n (a, b)
zip = (a -> b -> (a, b)) -> Vec n a -> Vec n b -> Vec n (a, b)
forall a b c (n :: Nat).
(a -> b -> c) -> Vec n a -> Vec n b -> Vec n c
zipWith (,)

findIndex :: (a -> Bool) -> Vec n a -> Maybe Int
findIndex :: (a -> Bool) -> Vec n a -> Maybe Int
findIndex = Int -> (a -> Bool) -> Vec n a -> Maybe Int
forall a (n :: Nat). Int -> (a -> Bool) -> Vec n a -> Maybe Int
go Int
0
  where
    go :: Int -> (a -> Bool) -> Vec n a -> Maybe Int
    go :: Int -> (a -> Bool) -> Vec n a -> Maybe Int
go Int
_ a -> Bool
_ Vec n a
Nil = Maybe Int
forall a. Maybe a
Nothing
    go Int
acc a -> Bool
p (Cons a
x Vec n a
xs)
      | a -> Bool
p a
x = Int -> Maybe Int
forall a. a -> Maybe a
Just Int
acc
      | Bool
otherwise = Int -> (a -> Bool) -> Vec n a -> Maybe Int
forall a (n :: Nat). Int -> (a -> Bool) -> Vec n a -> Maybe Int
go (Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) a -> Bool
p Vec n a
xs

(!!) :: Vec n a -> Int -> a
Cons a
x Vec n a
_  !! :: Vec n a -> Int -> a
!! Int
0 = a
x
Cons a
_ Vec n a
v' !! Int
n = Vec n a
v' Vec n a -> Int -> a
forall (n :: Nat) a. Vec n a -> Int -> a
!! (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
Vec n a
_ !! Int
_         = String -> a
forall a. HasCallStack => String -> a
error String
"Vec: (!!)"

replace :: Int -> a -> Vec n a -> Vec n a
replace :: Int -> a -> Vec n a -> Vec n a
replace Int
0 a
a (Cons a
_ Vec n a
xs) = a -> Vec n a -> Vec ('S n) a
forall a (n :: Nat). a -> Vec n a -> Vec ('S n) a
Cons a
a Vec n a
xs
replace Int
n a
a (Cons a
x Vec n a
xs) = a -> Vec n a -> Vec ('S n) a
forall a (n :: Nat). a -> Vec n a -> Vec ('S n) a
Cons a
x (Int -> a -> Vec n a -> Vec n a
forall a (n :: Nat). Int -> a -> Vec n a -> Vec n a
replace (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) a
a Vec n a
xs)
replace Int
_ a
_ Vec n a
Nil = String -> Vec n a
forall a. HasCallStack => String -> a
error String
"Found asymmetry is beyond the limits."

-- Equality type
data EqT a b where
  ReflEq :: EqT a a

data ExistsEqT t n where
     ExistsEqT :: EqT (t m) n -> ExistsEqT t n

-- Lists existentially quanitify over a vector's size : Exists n . Vec n a
data VecBox a where
     VecBox :: Vec n a -> VecBox a

fromList :: [a] -> VecBox a
fromList :: [a] -> VecBox a
fromList = (a -> VecBox a -> VecBox a) -> VecBox a -> [a] -> VecBox a
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\a
x (VecBox Vec n a
xs) -> Vec ('S n) a -> VecBox a
forall (n :: Nat) a. Vec n a -> VecBox a
VecBox (a -> Vec n a -> Vec ('S n) a
forall a (n :: Nat). a -> Vec n a -> Vec ('S n) a
Cons a
x Vec n a
xs)) (Vec 'Z a -> VecBox a
forall (n :: Nat) a. Vec n a -> VecBox a
VecBox Vec 'Z a
forall a. Vec 'Z a
Nil)

toList :: Vec n a -> [ a ]
toList :: Vec n a -> [a]
toList Vec n a
Nil = [ ]
toList (Cons a
x Vec n a
xs) = a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: Vec n a -> [a]
forall (n :: Nat) a. Vec n a -> [a]
toList Vec n a
xs

proveEqSize :: Vec n a -> Vec m b -> Maybe (EqT m n)
proveEqSize :: Vec n a -> Vec m b -> Maybe (EqT m n)
proveEqSize Vec n a
Nil Vec m b
Nil = EqT m m -> Maybe (EqT m m)
forall (m :: * -> *) a. Monad m => a -> m a
return EqT m m
forall k (a :: k). EqT a a
ReflEq
proveEqSize (Cons a
_ Vec n a
xs) (Cons b
_ Vec n b
ys) = do
  EqT n n
ReflEq <- Vec n a -> Vec n b -> Maybe (EqT n n)
forall (n :: Nat) a (m :: Nat) b.
Vec n a -> Vec m b -> Maybe (EqT m n)
proveEqSize Vec n a
xs Vec n b
ys
  EqT m m -> Maybe (EqT m m)
forall (m :: * -> *) a. Monad m => a -> m a
return EqT m m
forall k (a :: k). EqT a a
ReflEq
proveEqSize Vec n a
_ Vec m b
_ = Maybe (EqT m n)
forall a. Maybe a
Nothing

proveNonEmpty :: Vec n a -> Maybe (ExistsEqT 'S n)
proveNonEmpty :: Vec n a -> Maybe (ExistsEqT 'S n)
proveNonEmpty Vec n a
v =
  case Vec n a
v of
    Vec n a
Nil -> Maybe (ExistsEqT 'S n)
forall a. Maybe a
Nothing
    (Cons a
_ Vec n a
_) -> ExistsEqT 'S ('S n) -> Maybe (ExistsEqT 'S ('S n))
forall a. a -> Maybe a
Just (ExistsEqT 'S ('S n) -> Maybe (ExistsEqT 'S ('S n)))
-> ExistsEqT 'S ('S n) -> Maybe (ExistsEqT 'S ('S n))
forall a b. (a -> b) -> a -> b
$ EqT ('S n) ('S n) -> ExistsEqT 'S ('S n)
forall k k (t :: k -> k) (m :: k) (n :: k).
EqT (t m) n -> ExistsEqT t n
ExistsEqT EqT ('S n) ('S n)
forall k (a :: k). EqT a a
ReflEq

{- Vector list repreentation where the size 'n' is existential quantified -}
data VecList a where VL :: [Vec n a] -> VecList a

-- pre-condition: the input is a 'rectangular' list of lists (i.e. all internal
-- lists have the same size)
fromLists :: forall a . [[a]] -> VecList a
fromLists :: [[a]] -> VecList a
fromLists [] = [Vec 'Z a] -> VecList a
forall (n :: Nat) a. [Vec n a] -> VecList a
VL ([] :: [Vec 'Z a])
fromLists ([a]
xs:[[a]]
xss) = VecBox a -> VecList a -> VecList a
consList ([a] -> VecBox a
forall a. [a] -> VecBox a
fromList [a]
xs) ([[a]] -> VecList a
forall a. [[a]] -> VecList a
fromLists [[a]]
xss)
  where
    consList :: VecBox a -> VecList a -> VecList a
    consList :: VecBox a -> VecList a -> VecList a
consList (VecBox Vec n a
vec) (VL [])     = [Vec n a] -> VecList a
forall (n :: Nat) a. [Vec n a] -> VecList a
VL [Vec n a
vec]
    consList (VecBox Vec n a
vec) (VL [Vec n a]
xs') = -- Force the pre-condition equality
      case Vec n a -> [Vec n a] -> EqT n n
forall (n :: Nat) (n1 :: Nat) a'.
Vec n a' -> [Vec n1 a'] -> EqT n n1
preCondition Vec n a
vec [Vec n a]
xs' of
          EqT n n
ReflEq -> [Vec n a] -> VecList a
forall (n :: Nat) a. [Vec n a] -> VecList a
VL (Vec n a
vec Vec n a -> [Vec n a] -> [Vec n a]
forall a. a -> [a] -> [a]
: [Vec n a]
[Vec n a]
xs')
          where -- At the moment the pre-condition is 'assumed', and therefore
            -- force used unsafeCoerce: TODO, rewrite
            preCondition :: forall n n1 a' . Vec n a' -> [Vec n1 a'] -> EqT n n1
            preCondition :: Vec n a' -> [Vec n1 a'] -> EqT n n1
preCondition Vec n a'
_ [Vec n1 a']
_ = EqT Any Any -> EqT n n1
forall a b. a -> b
unsafeCoerce EqT Any Any
forall k (a :: k). EqT a a
ReflEq