{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE Safe #-}
-- | The representable closure (doc/synolon.md §3.1): the fixed bit width
--   of a Synolon type. Computed once, here, for the two consumers that must
--   agree — the machine lint, which rejects a binder, cell, port, or halt
--   answer whose type has no width, and the fold, which lays values out at
--   that width — so the two agree on which types have a width (a type in a
--   position only the fold sizes, such as a definition's codomain or a
--   primitive's instantiation, gets the same diagnostic there).
--
--   Widths: @Vec n τ@ is @n · |τ|@ (@n@ nat-closed); @Finite n@ is
--   @nbits n@; @Integer@ is 128; @Proxy@ is 0; a tuple is the sum of its
--   components; a declared datatype is @nbits (#constructors)@ plus the
--   widest constructor payload (the sum of its field widths at the
--   instantiation); a type variable is 0 (the fold only sees closed types);
--   a recursive datatype, an open width, and a function type have no width.
module ReWire.Synolon.Repr (DataEnv (..), dataEnv, Sizes, sizeOf, sizeOfM, isTupleCon) where

import ReWire.BitVector (nbits)
import ReWire.Eidos.Pretty ()
import ReWire.Eidos.Syntax
import ReWire.Eidos.Types (flattenArrow, flattenTyApp, evalNat, natNorm, substTv)
import ReWire.Pretty (prettyPrint)

import Control.Monad.State.Strict (StateT, evalStateT, get, modify, lift)
import Data.HashMap.Strict (HashMap)
import Data.HashSet (HashSet)
import Data.List (genericLength)
import Data.Text (Text)
import Numeric.Natural (Natural)

import qualified Data.HashMap.Strict as Map
import qualified Data.HashSet        as Set
import qualified Data.Text           as T
import qualified ReWire.Annotation   as Ann

-- | The datatype table sizing consults: each datatype's constructors in
--   declaration order (a constructor's tag is its index) and each
--   constructor's declared signature.
data DataEnv = DataEnv
      { DataEnv -> HashMap Text [Text]
deCtors   :: HashMap TyConId [DataConId]
      , DataEnv -> HashMap Text Sig
deCtorSig :: HashMap DataConId Sig
      }

dataEnv :: [DataDefn] -> DataEnv
dataEnv :: [DataDefn] -> DataEnv
dataEnv [DataDefn]
datas = DataEnv
      { deCtors :: HashMap Text [Text]
deCtors   = [(Text, [Text])] -> HashMap Text [Text]
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
Map.fromList [ (DataDefn -> Text
dataName DataDefn
d, [ Text
c | DataCon Annote
_ Text
c Sig
_ <- DataDefn -> [DataCon]
dataCons DataDefn
d ]) | DataDefn
d <- [DataDefn]
datas ]
      , deCtorSig :: HashMap Text Sig
deCtorSig = [(Text, Sig)] -> HashMap Text Sig
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
Map.fromList [ (Text
c, Sig
sig) | DataDefn
d <- [DataDefn]
datas, DataCon Annote
_ Text
c Sig
sig <- DataDefn -> [DataCon]
dataCons DataDefn
d ]
      }

-- | The sizing memo (types compare structurally; annotations are ignored).
type Sizes = HashMap Ty Natural

-- | The width of a type, or the reason it has none.
sizeOf :: DataEnv -> Ty -> Either Text Natural
sizeOf :: DataEnv -> Ty -> Either Text Natural
sizeOf DataEnv
de Ty
t = StateT Sizes (Either Text) Natural -> Sizes -> Either Text Natural
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT (DataEnv -> Ty -> StateT Sizes (Either Text) Natural
sizeOfM DataEnv
de Ty
t) Sizes
forall a. Monoid a => a
mempty

-- | 'sizeOf' over an explicit memo (the fold sizes many types per program).
sizeOfM :: DataEnv -> Ty -> StateT Sizes (Either Text) Natural
sizeOfM :: DataEnv -> Ty -> StateT Sizes (Either Text) Natural
sizeOfM DataEnv
de = Int -> HashSet Ty -> Ty -> StateT Sizes (Either Text) Natural
go Int
depth0 HashSet Ty
forall a. Monoid a => a
mempty
      where -- The visited set holds the datatype instances on the current
            -- descent: meeting one again is a recursive datatype. Types
            -- are compared after 'natNorm' (as the validator compares
            -- them), so an instance recursive through type arithmetic
            -- (@T (n + 0)@ under @T n@) is met again; the depth bound
            -- covers an unfolding that never revisits an instance.
            go :: Int -> HashSet Ty -> Ty -> StateT Sizes (Either Text) Natural
            go :: Int -> HashSet Ty -> Ty -> StateT Sizes (Either Text) Natural
go Int
depth HashSet Ty
visited (Ty -> Ty
natNorm -> Ty
t)
                  | Int
depth Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = Either Text Natural -> StateT Sizes (Either Text) Natural
forall (m :: * -> *) a. Monad m => m a -> StateT Sizes m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Either Text Natural -> StateT Sizes (Either Text) Natural)
-> Either Text Natural -> StateT Sizes (Either Text) Natural
forall a b. (a -> b) -> a -> b
$ Text -> Either Text Natural
forall a b. a -> Either a b
Left (Text -> Either Text Natural) -> Text -> Either Text Natural
forall a b. (a -> b) -> a -> b
$ Text
"datatype unfolding too deep (is the datatype recursive through its type arguments?): " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Ty -> Text
pp Ty
t
                  | Bool
otherwise  = do
                  m <- StateT Sizes (Either Text) Sizes
forall s (m :: * -> *). MonadState s m => m s
get
                  case Map.lookup t m of
                        Just Natural
sz -> Natural -> StateT Sizes (Either Text) Natural
forall a. a -> StateT Sizes (Either Text) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Natural
sz
                        Maybe Natural
Nothing -> do
                              sz <- case Ty -> (Ty, [Ty])
flattenTyApp Ty
t of
                                    (TyCon Annote
_ Text
"Vec", [Ty
n, Ty
te])
                                          | Just Natural
k <- Ty -> Maybe Natural
evalNat Ty
n -> (Natural
k Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
*) (Natural -> Natural)
-> StateT Sizes (Either Text) Natural
-> StateT Sizes (Either Text) Natural
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> HashSet Ty -> Ty -> StateT Sizes (Either Text) Natural
go Int
depth HashSet Ty
visited Ty
te
                                          | Bool
otherwise           -> Either Text Natural -> StateT Sizes (Either Text) Natural
forall (m :: * -> *) a. Monad m => m a -> StateT Sizes m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Either Text Natural -> StateT Sizes (Either Text) Natural)
-> Either Text Natural -> StateT Sizes (Either Text) Natural
forall a b. (a -> b) -> a -> b
$ Text -> Either Text Natural
forall a b. a -> Either a b
Left (Text -> Either Text Natural) -> Text -> Either Text Natural
forall a b. (a -> b) -> a -> b
$ Text
"can't determine the size of a Vec. (" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Ty -> Text
pp Ty
t Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
")"
                                    (TyCon Annote
_ Text
"Finite", [Ty
n])
                                          | Just Natural
k <- Ty -> Maybe Natural
evalNat Ty
n -> Natural -> StateT Sizes (Either Text) Natural
forall a. a -> StateT Sizes (Either Text) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Natural -> StateT Sizes (Either Text) Natural)
-> Natural -> StateT Sizes (Either Text) Natural
forall a b. (a -> b) -> a -> b
$ Natural -> Natural
nbits Natural
k
                                          | Bool
otherwise           -> Either Text Natural -> StateT Sizes (Either Text) Natural
forall (m :: * -> *) a. Monad m => m a -> StateT Sizes m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Either Text Natural -> StateT Sizes (Either Text) Natural)
-> Either Text Natural -> StateT Sizes (Either Text) Natural
forall a b. (a -> b) -> a -> b
$ Text -> Either Text Natural
forall a b. a -> Either a b
Left (Text -> Either Text Natural) -> Text -> Either Text Natural
forall a b. (a -> b) -> a -> b
$ Text
"can't determine the size of a Finite. (" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Ty -> Text
pp Ty
t Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
")"
                                    (TyCon Annote
_ Text
"Integer", [])     -> Natural -> StateT Sizes (Either Text) Natural
forall a. a -> StateT Sizes (Either Text) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Natural
128
                                    (TyCon Annote
_ Text
"Proxy", [Ty]
_)        -> Natural -> StateT Sizes (Either Text) Natural
forall a. a -> StateT Sizes (Either Text) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Natural
0
                                    (TyCon Annote
_ Text
c, [Ty]
args)
                                          | Text -> Bool
isTupleCon Text
c           -> [Natural] -> Natural
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Natural] -> Natural)
-> StateT Sizes (Either Text) [Natural]
-> StateT Sizes (Either Text) Natural
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Ty -> StateT Sizes (Either Text) Natural)
-> [Ty] -> StateT Sizes (Either Text) [Natural]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Int -> HashSet Ty -> Ty -> StateT Sizes (Either Text) Natural
go Int
depth HashSet Ty
visited) [Ty]
args
                                          | Ty
t Ty -> HashSet Ty -> Bool
forall a. (Eq a, Hashable a) => a -> HashSet a -> Bool
`Set.member` HashSet Ty
visited -> Either Text Natural -> StateT Sizes (Either Text) Natural
forall (m :: * -> *) a. Monad m => m a -> StateT Sizes m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Either Text Natural -> StateT Sizes (Either Text) Natural)
-> Either Text Natural -> StateT Sizes (Either Text) Natural
forall a b. (a -> b) -> a -> b
$ Text -> Either Text Natural
forall a b. a -> Either a b
Left (Text -> Either Text Natural) -> Text -> Either Text Natural
forall a b. (a -> b) -> a -> b
$ Text
"can't determine the size of a recursive datatype: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
c
                                          | Just [Text]
ctors <- Text -> HashMap Text [Text] -> Maybe [Text]
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
Map.lookup Text
c (DataEnv -> HashMap Text [Text]
deCtors DataEnv
de) -> do
                                                ws <- (Text -> StateT Sizes (Either Text) Natural)
-> [Text] -> StateT Sizes (Either Text) [Natural]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Int
-> HashSet Ty -> Ty -> Text -> StateT Sizes (Either Text) Natural
ctorWidth (Int
depth Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Ty -> HashSet Ty -> HashSet Ty
forall a. (Eq a, Hashable a) => a -> HashSet a -> HashSet a
Set.insert Ty
t HashSet Ty
visited) Ty
t) [Text]
ctors
                                                pure $ nbits (genericLength ctors) + maximum (0 : ws)
                                    (TyVarT {}, [Ty]
_)              -> Natural -> StateT Sizes (Either Text) Natural
forall a. a -> StateT Sizes (Either Text) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Natural
0
                                    (Ty, [Ty])
_                           -> Either Text Natural -> StateT Sizes (Either Text) Natural
forall (m :: * -> *) a. Monad m => m a -> StateT Sizes m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Either Text Natural -> StateT Sizes (Either Text) Natural)
-> Either Text Natural -> StateT Sizes (Either Text) Natural
forall a b. (a -> b) -> a -> b
$ Text -> Either Text Natural
forall a b. a -> Either a b
Left (Text -> Either Text Natural) -> Text -> Either Text Natural
forall a b. (a -> b) -> a -> b
$ Text
"couldn't calculate the size of a type: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Ty -> Text
pp Ty
t
                              modify $ Map.insert t sz
                              pure sz

            -- A constructor's payload width at a concrete instance of its
            -- datatype: the field types, with the datatype's parameters
            -- matched against the instance.
            ctorWidth :: Int -> HashSet Ty -> Ty -> DataConId -> StateT Sizes (Either Text) Natural
            ctorWidth :: Int
-> HashSet Ty -> Ty -> Text -> StateT Sizes (Either Text) Natural
ctorWidth Int
depth HashSet Ty
visited Ty
t Text
d = case Text -> HashMap Text Sig -> Maybe Sig
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
Map.lookup Text
d (DataEnv -> HashMap Text Sig
deCtorSig DataEnv
de) of
                  Just (Sig [TyVar]
_ Ty
ct) -> do
                        let ([Ty]
targs, Ty
tres) = Ty -> ([Ty], Ty)
flattenArrow Ty
ct
                        sub <- Either Text (HashMap TyVar Ty)
-> StateT Sizes (Either Text) (HashMap TyVar Ty)
forall (m :: * -> *) a. Monad m => m a -> StateT Sizes m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Either Text (HashMap TyVar Ty)
 -> StateT Sizes (Either Text) (HashMap TyVar Ty))
-> Either Text (HashMap TyVar Ty)
-> StateT Sizes (Either Text) (HashMap TyVar Ty)
forall a b. (a -> b) -> a -> b
$ Ty -> Ty -> Either Text (HashMap TyVar Ty)
matchTy Ty
tres Ty
t
                        sum <$> mapM (go depth visited . substTv sub) targs
                  Maybe Sig
Nothing         -> Natural -> StateT Sizes (Either Text) Natural
forall a. a -> StateT Sizes (Either Text) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Natural
0

            -- Nesting depth of datatype unfoldings, far beyond any real
            -- program's.
            depth0 :: Int
            depth0 :: Int
depth0 = Int
10000

            pp :: Ty -> Text
            pp :: Ty -> Text
pp = Ty -> Text
forall a. Pretty a => a -> Text
prettyPrint (Ty -> Text) -> (Ty -> Ty) -> Ty -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ty -> Ty
forall d. Data d => d -> d
Ann.unAnn

-- | Match a constructor's declared result type against a concrete instance.
matchTy :: Ty -> Ty -> Either Text (HashMap TyVar Ty)
matchTy :: Ty -> Ty -> Either Text (HashMap TyVar Ty)
matchTy (TyApp Annote
_ Ty
t1 Ty
t2) (TyApp Annote
_ Ty
t1' Ty
t2') = do
      s1 <- Ty -> Ty -> Either Text (HashMap TyVar Ty)
matchTy Ty
t1 Ty
t1'
      s2 <- matchTy t2 t2'
      if and (Map.intersectionWith (==) s1 s2)
            then pure $ Map.union s1 s2
            else Left "inconsistent assignment of a type variable in a constructor signature (rwc bug)"
matchTy (TyVarT Annote
_ TyVar
v) Ty
t = HashMap TyVar Ty -> Either Text (HashMap TyVar Ty)
forall a. a -> Either Text a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (HashMap TyVar Ty -> Either Text (HashMap TyVar Ty))
-> HashMap TyVar Ty -> Either Text (HashMap TyVar Ty)
forall a b. (a -> b) -> a -> b
$ TyVar -> Ty -> HashMap TyVar Ty
forall k v. Hashable k => k -> v -> HashMap k v
Map.singleton TyVar
v Ty
t
matchTy Ty
_ Ty
_            = HashMap TyVar Ty -> Either Text (HashMap TyVar Ty)
forall a. a -> Either Text a
forall (f :: * -> *) a. Applicative f => a -> f a
pure HashMap TyVar Ty
forall a. Monoid a => a
mempty

-- | The tuple type constructors, @()@ included (the zero-component tuple).
isTupleCon :: Text -> Bool
isTupleCon :: Text -> Bool
isTupleCon Text
c = Text -> Int
T.length Text
c Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
2 Bool -> Bool -> Bool
&& HasCallStack => Text -> Char
Text -> Char
T.head Text
c Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'(' Bool -> Bool -> Bool
&& HasCallStack => Text -> Char
Text -> Char
T.last Text
c Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
')'
            Bool -> Bool -> Bool
&& (Char -> Bool) -> Text -> Bool
T.all (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
',') (HasCallStack => Text -> Text
Text -> Text
T.init (Text -> Text) -> Text -> Text
forall a b. (a -> b) -> a -> b
$ HasCallStack => Text -> Text
Text -> Text
T.tail Text
c)