{-# LANGUAGE TypeFamilies #-}
{- |
This module provides an array shape type,
that allows to store elements from a container
while preserving the container structure.
-}
module Data.Array.Comfort.Container (
   C(..), EqShape(..), NFShape(..), Indexed(..),
   ) where

import qualified Data.Array.Comfort.Shape as Shape

import Control.DeepSeq (NFData, rnf)

import qualified Data.NonEmpty.Map as NonEmptyMap
import qualified Data.NonEmpty.Set as NonEmptySet
import qualified Data.NonEmpty.Class as NonEmptyC
import qualified Data.NonEmpty as NonEmpty
import qualified Data.Empty as Empty
import qualified Data.Map as Map
import qualified Data.Set as Set
-- import qualified Data.Complex as Complex
import Data.Map (Map)
import Data.Set (Set)
import Data.Foldable (Foldable)
import Data.Maybe (fromMaybe)
-- import Data.Complex (Complex((:+)))



class (Foldable f) => C f where
   data Shape f
   shapeSize :: Shape f -> Int
   fromList :: Shape f -> [a] -> f a
   toShape :: f a -> Shape f

class (C f) => NFShape f where
   rnfShape :: Shape f -> ()

class (C f) => EqShape f where
   eqShape :: Shape f -> Shape f -> Bool

class (C f) => Indexed f where
   type Index f
   indices :: Shape f -> [Index f]
   unifiedSizeOffset ::
      (Shape.Checking check) =>
      Shape f -> (Int, Index f -> Shape.Result check Int)


instance (NFShape f) => NFData (Shape f) where
   rnf :: Shape f -> ()
rnf = forall (f :: * -> *). NFShape f => Shape f -> ()
rnfShape

instance (EqShape f) => Eq (Shape f) where
   == :: Shape f -> Shape f -> Bool
(==) = forall (f :: * -> *). EqShape f => Shape f -> Shape f -> Bool
eqShape

instance (C f) => Shape.C (Shape f) where
   size :: Shape f -> Int
size = forall (f :: * -> *). C f => Shape f -> Int
shapeSize

instance (Indexed f) => Shape.Indexed (Shape f) where
   type Index (Shape f) = Index f
   indices :: Shape f -> [Index (Shape f)]
indices = forall (f :: * -> *). Indexed f => Shape f -> [Index f]
indices
   unifiedSizeOffset :: forall check.
Checking check =>
Shape f -> (Int, Index (Shape f) -> Result check Int)
unifiedSizeOffset = forall (f :: * -> *) check.
(Indexed f, Checking check) =>
Shape f -> (Int, Index f -> Result check Int)
unifiedSizeOffset


instance C [] where
   data Shape [] = ShapeList Int
      deriving (Int -> Shape [] -> ShowS
[Shape []] -> ShowS
Shape [] -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Shape []] -> ShowS
$cshowList :: [Shape []] -> ShowS
show :: Shape [] -> String
$cshow :: Shape [] -> String
showsPrec :: Int -> Shape [] -> ShowS
$cshowsPrec :: Int -> Shape [] -> ShowS
Show)
   shapeSize :: Shape [] -> Int
shapeSize (ShapeList Int
n) = Int
n
   toShape :: forall a. [a] -> Shape []
toShape = Int -> Shape []
ShapeList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a. Foldable t => t a -> Int
length
   fromList :: forall a. Shape [] -> [a] -> [a]
fromList Shape []
_ = forall a. a -> a
id

instance EqShape [] where
   eqShape :: Shape [] -> Shape [] -> Bool
eqShape (ShapeList Int
n)  (ShapeList Int
m) = Int
nforall a. Eq a => a -> a -> Bool
==Int
m

instance NFShape [] where
   rnfShape :: Shape [] -> ()
rnfShape (ShapeList Int
n) = forall a. NFData a => a -> ()
rnf Int
n

instance Indexed [] where
   type Index [] = Int
   indices :: Shape [] -> [Index []]
indices (ShapeList Int
len) = forall a. Int -> [a] -> [a]
take Int
len forall a b. (a -> b) -> a -> b
$ forall a. (a -> a) -> a -> [a]
iterate (Int
1forall a. Num a => a -> a -> a
+) Int
0
   unifiedSizeOffset :: forall check.
Checking check =>
Shape [] -> (Int, Index [] -> Result check Int)
unifiedSizeOffset (ShapeList Int
len) =
      (Int
len, \Index []
ix -> do
         forall check. Checking check => String -> Bool -> Result check ()
Shape.assert String
"Shape.Container.[]: array index too small" forall a b. (a -> b) -> a -> b
$ Index []
ixforall a. Ord a => a -> a -> Bool
>=Int
0
         forall check. Checking check => String -> Bool -> Result check ()
Shape.assert String
"Shape.Container.[]: array index too big" forall a b. (a -> b) -> a -> b
$ Index []
ixforall a. Ord a => a -> a -> Bool
<Int
len
         forall (m :: * -> *) a. Monad m => a -> m a
return Index []
ix)


{-
instance Foldable only available since GHC-8.0. :-(
Could be circumvented by Data.Orphans
but that one also pulls in lots of dangerous instances.

instance C Complex where
   data Shape Complex = ShapeComplex
   shapeSize ShapeComplex = 2
   toShape (_:+_) = ShapeComplex
   fromList ShapeComplex xs =
      case xs of
         [r,i] -> r Complex.:+ i
         _ -> error "ShapeComplex: not two elements"

instance EqShape Complex where
   eqShape ShapeComplex ShapeComplex = True

instance NFShape Complex where
   rnfShape ShapeComplex = ()
-}


instance (C f) => C (NonEmpty.T f) where
   data Shape (NonEmpty.T f) = ShapeNonEmpty (Shape f)
   shapeSize :: Shape (T f) -> Int
shapeSize (ShapeNonEmpty Shape f
c) = Int
1 forall a. Num a => a -> a -> a
+ forall (f :: * -> *). C f => Shape f -> Int
shapeSize Shape f
c
   toShape :: forall a. T f a -> Shape (T f)
toShape = forall (f :: * -> *). Shape f -> Shape (T f)
ShapeNonEmpty forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a. C f => f a -> Shape f
toShape forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a. T f a -> f a
NonEmpty.tail
   fromList :: forall a. Shape (T f) -> [a] -> T f a
fromList (ShapeNonEmpty Shape f
c) [a]
xt =
      case [a]
xt of
         [] -> forall a. HasCallStack => String -> a
error String
"ShapeNonEmpty: empty list"
         a
x:[a]
xs -> forall a (f :: * -> *). a -> f a -> T f a
NonEmpty.cons a
x forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. C f => Shape f -> [a] -> f a
fromList Shape f
c [a]
xs

instance (EqShape f) => EqShape (NonEmpty.T f) where
   eqShape :: Shape (T f) -> Shape (T f) -> Bool
eqShape (ShapeNonEmpty Shape f
a) (ShapeNonEmpty Shape f
b) = Shape f
aforall a. Eq a => a -> a -> Bool
==Shape f
b

instance (NFShape f) => NFShape (NonEmpty.T f) where
   rnfShape :: Shape (T f) -> ()
rnfShape (ShapeNonEmpty Shape f
c) = forall (f :: * -> *). NFShape f => Shape f -> ()
rnfShape Shape f
c

instance (C f) => Indexed (NonEmpty.T f) where
   type Index (NonEmpty.T f) = Int
   indices :: Shape (T f) -> [Index (T f)]
indices Shape (T f)
shape = forall a. Int -> [a] -> [a]
take (forall (f :: * -> *). C f => Shape f -> Int
shapeSize Shape (T f)
shape) forall a b. (a -> b) -> a -> b
$ forall a. (a -> a) -> a -> [a]
iterate (Int
1forall a. Num a => a -> a -> a
+) Int
0
   unifiedSizeOffset :: forall check.
Checking check =>
Shape (T f) -> (Int, Index (T f) -> Result check Int)
unifiedSizeOffset Shape (T f)
shape =
      let len :: Int
len = forall (f :: * -> *). C f => Shape f -> Int
shapeSize Shape (T f)
shape in
      (Int
len, \Index (T f)
ix -> do
         forall check. Checking check => String -> Bool -> Result check ()
Shape.assert String
"Shape.Container.NonEmpty: array index too small" forall a b. (a -> b) -> a -> b
$ Index (T f)
ixforall a. Ord a => a -> a -> Bool
>=Int
0
         forall check. Checking check => String -> Bool -> Result check ()
Shape.assert String
"Shape.Container.NonEmpty: array index too big" forall a b. (a -> b) -> a -> b
$ Index (T f)
ixforall a. Ord a => a -> a -> Bool
<Int
len
         forall (m :: * -> *) a. Monad m => a -> m a
return Index (T f)
ix)


instance C Empty.T where
   data Shape Empty.T = ShapeEmpty
      deriving (Int -> Shape T -> ShowS
[Shape T] -> ShowS
Shape T -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Shape T] -> ShowS
$cshowList :: [Shape T] -> ShowS
show :: Shape T -> String
$cshow :: Shape T -> String
showsPrec :: Int -> Shape T -> ShowS
$cshowsPrec :: Int -> Shape T -> ShowS
Show)
   shapeSize :: Shape T -> Int
shapeSize Shape T
R:ShapeT1
ShapeEmpty = Int
0
   toShape :: forall a. T a -> Shape T
toShape T a
Empty.Cons = Shape T
ShapeEmpty
   fromList :: forall a. Shape T -> [a] -> T a
fromList Shape T
R:ShapeT1
ShapeEmpty [a]
xs =
      case [a]
xs of
         [] -> forall a. T a
Empty.Cons
         [a]
_ -> forall a. HasCallStack => String -> a
error String
"ShapeEmpty: not empty"

instance EqShape Empty.T where
   eqShape :: Shape T -> Shape T -> Bool
eqShape Shape T
R:ShapeT1
ShapeEmpty Shape T
R:ShapeT1
ShapeEmpty = Bool
True

instance NFShape Empty.T where
   rnfShape :: Shape T -> ()
rnfShape Shape T
R:ShapeT1
ShapeEmpty = ()


instance (Ord k) => C (Map k) where
   data Shape (Map k) = ShapeMap (Set k)
      deriving (Int -> Shape (Map k) -> ShowS
forall k. Show k => Int -> Shape (Map k) -> ShowS
forall k. Show k => [Shape (Map k)] -> ShowS
forall k. Show k => Shape (Map k) -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Shape (Map k)] -> ShowS
$cshowList :: forall k. Show k => [Shape (Map k)] -> ShowS
show :: Shape (Map k) -> String
$cshow :: forall k. Show k => Shape (Map k) -> String
showsPrec :: Int -> Shape (Map k) -> ShowS
$cshowsPrec :: forall k. Show k => Int -> Shape (Map k) -> ShowS
Show)
   shapeSize :: Shape (Map k) -> Int
shapeSize (ShapeMap Set k
set) = forall a. Set a -> Int
Set.size Set k
set
   toShape :: forall a. Map k a -> Shape (Map k)
toShape = forall k. Set k -> Shape (Map k)
ShapeMap forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Map k a -> Set k
Map.keysSet
   fromList :: forall a. Shape (Map k) -> [a] -> Map k a
fromList (ShapeMap Set k
set) = forall k a. Eq k => [(k, a)] -> Map k a
Map.fromAscList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. [a] -> [b] -> [(a, b)]
zip (forall a. Set a -> [a]
Set.toAscList Set k
set)

instance (Ord k) => EqShape (Map k) where
   eqShape :: Shape (Map k) -> Shape (Map k) -> Bool
eqShape (ShapeMap Set k
set0) (ShapeMap Set k
set1) = Set k
set0forall a. Eq a => a -> a -> Bool
==Set k
set1

instance (NFData k, Ord k) => NFShape (Map k) where
   rnfShape :: Shape (Map k) -> ()
rnfShape (ShapeMap Set k
set) = forall a. NFData a => a -> ()
rnf Set k
set

instance (Ord k) => Indexed (Map k) where
   type Index (Map k) = k
   indices :: Shape (Map k) -> [Index (Map k)]
indices (ShapeMap Set k
set) = forall a. Set a -> [a]
Set.toAscList Set k
set
   unifiedSizeOffset :: forall check.
Checking check =>
Shape (Map k) -> (Int, Index (Map k) -> Result check Int)
unifiedSizeOffset (ShapeMap Set k
set) = forall sh check.
(Indexed sh, Checking check) =>
sh -> (Int, Index sh -> Result check Int)
Shape.unifiedSizeOffset Set k
set


instance (Ord k) => C (NonEmptyMap.T k) where
   data Shape (NonEmptyMap.T k) = ShapeNonEmptyMap (NonEmptySet.T k)
      deriving (Int -> Shape (T k) -> ShowS
forall k. Show k => Int -> Shape (T k) -> ShowS
forall k. Show k => [Shape (T k)] -> ShowS
forall k. Show k => Shape (T k) -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Shape (T k)] -> ShowS
$cshowList :: forall k. Show k => [Shape (T k)] -> ShowS
show :: Shape (T k) -> String
$cshow :: forall k. Show k => Shape (T k) -> String
showsPrec :: Int -> Shape (T k) -> ShowS
$cshowsPrec :: forall k. Show k => Int -> Shape (T k) -> ShowS
Show)
   shapeSize :: Shape (T k) -> Int
shapeSize (ShapeNonEmptyMap T k
set) = forall a. T a -> Int
NonEmptySet.size T k
set
   toShape :: forall a. T k a -> Shape (T k)
toShape = forall k. T k -> Shape (T k)
ShapeNonEmptyMap forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => T k a -> T k
NonEmptyMap.keysSet
   fromList :: forall a. Shape (T k) -> [a] -> T k a
fromList (ShapeNonEmptyMap T k
set) =
      forall k a. Ord k => T [] (k, a) -> T k a
NonEmptyMap.fromAscList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Zip f => f a -> f b -> f (a, b)
NonEmptyC.zip (forall a. T a -> T [] a
NonEmptySet.toAscList T k
set) forall b c a. (b -> c) -> (a -> b) -> a -> c
.
      forall a. a -> Maybe a -> a
fromMaybe (forall a. HasCallStack => String -> a
error String
"ShapeNonEmptyMap: empty list") forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a. ViewL f => f a -> Maybe (T f a)
NonEmpty.fetch

instance (Ord k) => EqShape (NonEmptyMap.T k) where
   eqShape :: Shape (T k) -> Shape (T k) -> Bool
eqShape (ShapeNonEmptyMap T k
set0) (ShapeNonEmptyMap T k
set1) = T k
set0forall a. Eq a => a -> a -> Bool
==T k
set1

instance (NFData k, Ord k) => NFShape (NonEmptyMap.T k) where
   rnfShape :: Shape (T k) -> ()
rnfShape (ShapeNonEmptyMap T k
set) = forall a. NFData a => a -> ()
rnf T k
set

instance (Ord k) => Indexed (NonEmptyMap.T k) where
   type Index (NonEmptyMap.T k) = k
   indices :: Shape (T k) -> [Index (T k)]
indices (ShapeNonEmptyMap T k
set) =
      forall (f :: * -> *) a. Cons f => T f a -> f a
NonEmpty.flatten forall a b. (a -> b) -> a -> b
$ forall a. T a -> T [] a
NonEmptySet.toAscList T k
set
   unifiedSizeOffset :: forall check.
Checking check =>
Shape (T k) -> (Int, Index (T k) -> Result check Int)
unifiedSizeOffset (ShapeNonEmptyMap T k
set) =
      forall sh check.
(Indexed sh, Checking check) =>
sh -> (Int, Index sh -> Result check Int)
Shape.unifiedSizeOffset (forall a. Ord a => T a -> Set a
NonEmptySet.flatten T k
set)