```-- |
-- Module      :  Cryptol.Transform.MonoValues
-- Copyright   :  (c) 2013-2016 Galois, Inc.
-- Maintainer  :  cryptol@galois.com
-- Stability   :  provisional
-- Portability :  portable
--
-- This module implements a transformation, which tries to avoid exponential
-- slow down in some cases.  What's the problem?  Consider the following (common)
-- patterns:
--
-- >    fibs = [0,1] # [ x + y | x <- fibs, y <- drop`{1} fibs ]
--
-- The type of @fibs@ is:
--
-- >    {a} (a >= 1, fin a) => [inf][a]
--
-- Here @a@ is the number of bits to be used in the values computed by @fibs@.
-- When we evaluate @fibs@, @a@ becomes a parameter to @fibs@, which works
-- except that now @fibs@ is a function, and we don't get any of the memoization
-- we might expect!  What looked like an efficient implementation has all
-- of a sudden become exponential!
--
-- Note that this is only a problem for polymorphic values: if @fibs@ was
-- already a function, it would not be that surprising that it does not
-- get cached.
--
-- So, to avoid this, we try to spot recursive polymorphic values,
-- where the recursive occurrences have the exact same type parameters
-- as the definition.  For example, this is the case in @fibs@: each
-- recursive call to @fibs@ is instantiated with exactly the same
-- type parameter (i.e., @a@).  The rewrite we do is as follows:
--
-- >    fibs : {a} (a >= 1, fin a) => [inf][a]
-- >    fibs = \{a} (a >= 1, fin a) -> fibs'
-- >      where fibs' : [inf][a]
-- >            fibs' = [0,1] # [ x + y | x <- fibs', y <- drop`{1} fibs' ]
--
-- After the rewrite, the recursion is monomorphic (i.e., we are always using
-- the same type).  As a result, @fibs'@ is an ordinary recursive value,
-- where we get the benefit of caching.
--
-- The rewrite is a bit more complex, when there are multiple mutually
-- recursive functions.  Here is an example:
--
-- >    zig : {a} (a >= 2, fin a) => [inf][a]
-- >    zig = [1] # zag
-- >
-- >    zag : {a} (a >= 2, fin a) => [inf][a]
-- >    zag = [2] # zig
--
-- This gets rewritten to:
--
-- >    newName : {a} (a >= 2, fin a) => ([inf][a], [inf][a])
-- >    newName = \{a} (a >= 2, fin a) -> (zig', zag')
-- >      where
-- >      zig' : [inf][a]
-- >      zig' = [1] # zag'
-- >
-- >      zag' : [inf][a]
-- >      zag' = [2] # zig'
-- >
-- >    zig : {a} (a >= 2, fin a) => [inf][a]
-- >    zig = \{a} (a >= 2, fin a) -> (newName a <> <> ).1
-- >
-- >    zag : {a} (a >= 2, fin a) => [inf][a]
-- >    zag = \{a} (a >= 2, fin a) -> (newName a <> <> ).2
--
-- NOTE:  We are assuming that no capture would occur with binders.
-- For values, this is because we replaces things with freshly chosen variables.
-- For types, this should be because there should be no shadowing in the types.
-- XXX: Make sure that this really is the case for types!!

{-# LANGUAGE PatternGuards, FlexibleInstances, MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE CPP #-}
module Cryptol.Transform.MonoValues (rewModule) where

import Cryptol.ModuleSystem.Name
(SupplyT,liftSupply,Supply,mkDeclared,NameSource(..))
import Cryptol.Parser.Position (emptyRange)
import Cryptol.TypeCheck.AST hiding (splitTApp) -- XXX: just use this one
import Cryptol.TypeCheck.TypeMap
import Cryptol.Utils.Ident (ModName)
import Data.List(sortBy,groupBy)
import Data.Either(partitionEithers)
import Data.Map (Map)

import Prelude ()
import Prelude.Compat

{- (f,t,n) |--> x  means that when we spot instantiations of @f@ with @ts@ and
@n@ proof argument, we should replace them with @Var x@ -}
newtype RewMap' a = RM (Map Name (TypesMap (Map Int a)))
type RewMap = RewMap' Name

instance TrieMap RewMap' (Name,[Type],Int) where
emptyTM :: RewMap' a
emptyTM  = Map Name (TypesMap (Map Int a)) -> RewMap' a
forall a. Map Name (TypesMap (Map Int a)) -> RewMap' a
RM Map Name (TypesMap (Map Int a))
forall (m :: * -> *) k a. TrieMap m k => m a
emptyTM

nullTM :: RewMap' a -> Bool
nullTM (RM Map Name (TypesMap (Map Int a))
m) = Map Name (TypesMap (Map Int a)) -> Bool
forall (m :: * -> *) k a. TrieMap m k => m a -> Bool
nullTM Map Name (TypesMap (Map Int a))
m

lookupTM :: (Name, [Type], Int) -> RewMap' a -> Maybe a
lookupTM (Name
x,[Type]
ts,Int
n) (RM Map Name (TypesMap (Map Int a))
m) = do TypesMap (Map Int a)
tM <- Name
-> Map Name (TypesMap (Map Int a)) -> Maybe (TypesMap (Map Int a))
forall (m :: * -> *) k a. TrieMap m k => k -> m a -> Maybe a
lookupTM Name
x Map Name (TypesMap (Map Int a))
m
Map Int a
tP <- [Type] -> TypesMap (Map Int a) -> Maybe (Map Int a)
forall (m :: * -> *) k a. TrieMap m k => k -> m a -> Maybe a
lookupTM [Type]
ts TypesMap (Map Int a)
tM
Int -> Map Int a -> Maybe a
forall (m :: * -> *) k a. TrieMap m k => k -> m a -> Maybe a
lookupTM Int
n Map Int a
tP

alterTM :: (Name, [Type], Int)
-> (Maybe a -> Maybe a) -> RewMap' a -> RewMap' a
alterTM (Name
x,[Type]
ts,Int
n) Maybe a -> Maybe a
f (RM Map Name (TypesMap (Map Int a))
m) = Map Name (TypesMap (Map Int a)) -> RewMap' a
forall a. Map Name (TypesMap (Map Int a)) -> RewMap' a
RM (Name
-> (Maybe (TypesMap (Map Int a)) -> Maybe (TypesMap (Map Int a)))
-> Map Name (TypesMap (Map Int a))
-> Map Name (TypesMap (Map Int a))
forall (m :: * -> *) k a.
TrieMap m k =>
k -> (Maybe a -> Maybe a) -> m a -> m a
alterTM Name
x Maybe (TypesMap (Map Int a)) -> Maybe (TypesMap (Map Int a))
forall (m :: * -> *) (m :: * -> *).
(TrieMap m [Type], TrieMap m Int) =>
Maybe (m (m a)) -> Maybe (m (m a))
f1 Map Name (TypesMap (Map Int a))
m)
where
f1 :: Maybe (m (m a)) -> Maybe (m (m a))
f1 Maybe (m (m a))
Nothing   = do a
a <- Maybe a -> Maybe a
f Maybe a
forall a. Maybe a
Nothing
m (m a) -> Maybe (m (m a))
forall (m :: * -> *) a. Monad m => a -> m a
return ([Type] -> m a -> m (m a) -> m (m a)
forall (m :: * -> *) k a. TrieMap m k => k -> a -> m a -> m a
insertTM [Type]
ts (Int -> a -> m a -> m a
forall (m :: * -> *) k a. TrieMap m k => k -> a -> m a -> m a
insertTM Int
n a
a m a
forall (m :: * -> *) k a. TrieMap m k => m a
emptyTM) m (m a)
forall (m :: * -> *) k a. TrieMap m k => m a
emptyTM)

f1 (Just m (m a)
tM) = m (m a) -> Maybe (m (m a))
forall a. a -> Maybe a
Just ([Type] -> (Maybe (m a) -> Maybe (m a)) -> m (m a) -> m (m a)
forall (m :: * -> *) k a.
TrieMap m k =>
k -> (Maybe a -> Maybe a) -> m a -> m a
alterTM [Type]
ts Maybe (m a) -> Maybe (m a)
forall (m :: * -> *). TrieMap m Int => Maybe (m a) -> Maybe (m a)
f2 m (m a)
tM)

f2 :: Maybe (m a) -> Maybe (m a)
f2 Maybe (m a)
Nothing   = do a
a <- Maybe a -> Maybe a
f Maybe a
forall a. Maybe a
Nothing
m a -> Maybe (m a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> a -> m a -> m a
forall (m :: * -> *) k a. TrieMap m k => k -> a -> m a -> m a
insertTM Int
n a
a m a
forall (m :: * -> *) k a. TrieMap m k => m a
emptyTM)

f2 (Just m a
pM) = m a -> Maybe (m a)
forall a. a -> Maybe a
Just (Int -> (Maybe a -> Maybe a) -> m a -> m a
forall (m :: * -> *) k a.
TrieMap m k =>
k -> (Maybe a -> Maybe a) -> m a -> m a
alterTM Int
n Maybe a -> Maybe a
f m a
pM)

unionTM :: (a -> a -> a) -> RewMap' a -> RewMap' a -> RewMap' a
unionTM a -> a -> a
f (RM Map Name (TypesMap (Map Int a))
a) (RM Map Name (TypesMap (Map Int a))
b) = Map Name (TypesMap (Map Int a)) -> RewMap' a
forall a. Map Name (TypesMap (Map Int a)) -> RewMap' a
RM ((TypesMap (Map Int a)
-> TypesMap (Map Int a) -> TypesMap (Map Int a))
-> Map Name (TypesMap (Map Int a))
-> Map Name (TypesMap (Map Int a))
-> Map Name (TypesMap (Map Int a))
forall (m :: * -> *) k a.
TrieMap m k =>
(a -> a -> a) -> m a -> m a -> m a
unionTM ((Map Int a -> Map Int a -> Map Int a)
-> TypesMap (Map Int a)
-> TypesMap (Map Int a)
-> TypesMap (Map Int a)
forall (m :: * -> *) k a.
TrieMap m k =>
(a -> a -> a) -> m a -> m a -> m a
unionTM ((a -> a -> a) -> Map Int a -> Map Int a -> Map Int a
forall (m :: * -> *) k a.
TrieMap m k =>
(a -> a -> a) -> m a -> m a -> m a
unionTM a -> a -> a
f)) Map Name (TypesMap (Map Int a))
a Map Name (TypesMap (Map Int a))
b)

toListTM :: RewMap' a -> [((Name, [Type], Int), a)]
toListTM (RM Map Name (TypesMap (Map Int a))
m) = [ ((Name
x,[Type]
ts,Int
n),a
y) | (Name
x,TypesMap (Map Int a)
tM)  <- Map Name (TypesMap (Map Int a)) -> [(Name, TypesMap (Map Int a))]
forall (m :: * -> *) k a. TrieMap m k => m a -> [(k, a)]
toListTM Map Name (TypesMap (Map Int a))
m
, ([Type]
ts,Map Int a
pM) <- TypesMap (Map Int a) -> [([Type], Map Int a)]
forall (m :: * -> *) k a. TrieMap m k => m a -> [(k, a)]
toListTM TypesMap (Map Int a)
tM
, (Int
n,a
y)   <- Map Int a -> [(Int, a)]
forall (m :: * -> *) k a. TrieMap m k => m a -> [(k, a)]
toListTM Map Int a
pM ]

mapMaybeWithKeyTM :: ((Name, [Type], Int) -> a -> Maybe b) -> RewMap' a -> RewMap' b
mapMaybeWithKeyTM (Name, [Type], Int) -> a -> Maybe b
f (RM Map Name (TypesMap (Map Int a))
m) =
Map Name (TypesMap (Map Int b)) -> RewMap' b
forall a. Map Name (TypesMap (Map Int a)) -> RewMap' a
RM ((Name -> TypesMap (Map Int a) -> TypesMap (Map Int b))
-> Map Name (TypesMap (Map Int a))
-> Map Name (TypesMap (Map Int b))
forall (m :: * -> *) k a b.
TrieMap m k =>
(k -> a -> b) -> m a -> m b
mapWithKeyTM      (\Name
qn  TypesMap (Map Int a)
tm ->
([Type] -> Map Int a -> Map Int b)
-> TypesMap (Map Int a) -> TypesMap (Map Int b)
forall (m :: * -> *) k a b.
TrieMap m k =>
(k -> a -> b) -> m a -> m b
mapWithKeyTM      (\[Type]
tys Map Int a
is ->
(Int -> a -> Maybe b) -> Map Int a -> Map Int b
forall (m :: * -> *) k a b.
TrieMap m k =>
(k -> a -> Maybe b) -> m a -> m b
mapMaybeWithKeyTM (\Int
i   a
a  -> (Name, [Type], Int) -> a -> Maybe b
f (Name
qn,[Type]
tys,Int
i) a
a) Map Int a
is) TypesMap (Map Int a)
tm) Map Name (TypesMap (Map Int a))
m)

-- | Note that this assumes that this pass will be run only once for each
-- module, otherwise we will get name collisions.
rewModule :: Supply -> Module -> (Module,Supply)
rewModule :: Supply -> Module -> (Module, Supply)
rewModule Supply
s Module
m = ReaderT RO (SupplyT Id) Module -> RO -> Supply -> (Module, Supply)
forall (m :: * -> *) a r. RunM m a r => m a -> r
runM ReaderT RO (SupplyT Id) Module
body (Module -> RO
mName Module
m) Supply
s
where
body :: ReaderT RO (SupplyT Id) Module
body = do [DeclGroup]
ds <- (DeclGroup -> ReaderT RO (SupplyT Id) DeclGroup)
-> [DeclGroup] -> ReaderT RO (SupplyT Id) [DeclGroup]
forall (t :: * -> *) (m :: * -> *) a b.
(a -> m b) -> t a -> m (t b)
mapM (RewMap -> DeclGroup -> ReaderT RO (SupplyT Id) DeclGroup
rewDeclGroup RewMap
forall (m :: * -> *) k a. TrieMap m k => m a
emptyTM) (Module -> [DeclGroup]
mDecls Module
m)
Module -> ReaderT RO (SupplyT Id) Module
forall (m :: * -> *) a. Monad m => a -> m a
return Module
m { mDecls :: [DeclGroup]
mDecls = [DeclGroup]
ds }

--------------------------------------------------------------------------------

type M  = ReaderT RO (SupplyT Id)
type RO = ModName

-- | Produce a fresh top-level name.
newName :: M Name
newName :: M Name
newName  =
do RO
ns <- ReaderT RO (SupplyT Id) RO
forall (m :: * -> *) i. ReaderM m i => m i
(Supply -> (Name, Supply)) -> M Name
forall (m :: * -> *) a. FreshM m => (Supply -> (a, Supply)) -> m a
liftSupply (RO
-> NameSource
-> Ident
-> Maybe Fixity
-> Range
-> Supply
-> (Name, Supply)
mkDeclared RO
ns NameSource
SystemName Ident
"\$mono" Maybe Fixity
forall a. Maybe a
Nothing Range
emptyRange)

newTopOrLocalName :: M Name
newTopOrLocalName :: M Name
newTopOrLocalName  = M Name
newName

-- | Not really any distinction between global and local, all names get the
-- module prefix added, and a unique id.
inLocal :: M a -> M a
inLocal :: M a -> M a
inLocal  = M a -> M a
forall a. a -> a
id

--------------------------------------------------------------------------------
rewE :: RewMap -> Expr -> M Expr   -- XXX: not IO
rewE :: RewMap -> Expr -> M Expr
rewE RewMap
rews = Expr -> M Expr
go

where
tryRewrite :: (Expr, [Type], Int) -> Maybe Expr
tryRewrite (EVar Name
x,[Type]
tps,Int
n) =
do Name
y <- (Name, [Type], Int) -> RewMap -> Maybe Name
forall (m :: * -> *) k a. TrieMap m k => k -> m a -> Maybe a
lookupTM (Name
x,[Type]
tps,Int
n) RewMap
rews
Expr -> Maybe Expr
forall (m :: * -> *) a. Monad m => a -> m a
return (Name -> Expr
EVar Name
y)
tryRewrite (Expr, [Type], Int)
_ = Maybe Expr
forall a. Maybe a
Nothing

go :: Expr -> M Expr
go Expr
expr =
case Expr
expr of

-- Interesting cases
ETApp Expr
e Type
t      -> case (Expr, [Type], Int) -> Maybe Expr
tryRewrite (Expr -> Int -> (Expr, [Type], Int)
splitTApp Expr
expr Int
0) of
Maybe Expr
Nothing  -> Expr -> Type -> Expr
ETApp (Expr -> Type -> Expr)
-> M Expr -> ReaderT RO (SupplyT Id) (Type -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<\$> Expr -> M Expr
go Expr
e ReaderT RO (SupplyT Id) (Type -> Expr)
-> ReaderT RO (SupplyT Id) Type -> M Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Type -> ReaderT RO (SupplyT Id) Type
forall (m :: * -> *) a. Monad m => a -> m a
return Type
t
Just Expr
yes -> Expr -> M Expr
forall (m :: * -> *) a. Monad m => a -> m a
return Expr
yes
EProofApp Expr
e    -> case (Expr, [Type], Int) -> Maybe Expr
tryRewrite (Expr -> Int -> (Expr, [Type], Int)
splitTApp Expr
e Int
1) of
Maybe Expr
Nothing  -> Expr -> Expr
EProofApp (Expr -> Expr) -> M Expr -> M Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<\$> Expr -> M Expr
go Expr
e
Just Expr
yes -> Expr -> M Expr
forall (m :: * -> *) a. Monad m => a -> m a
return Expr
yes

EList [Expr]
es Type
t      -> [Expr] -> Type -> Expr
EList   ([Expr] -> Type -> Expr)
-> ReaderT RO (SupplyT Id) [Expr]
-> ReaderT RO (SupplyT Id) (Type -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<\$> (Expr -> M Expr) -> [Expr] -> ReaderT RO (SupplyT Id) [Expr]
forall (t :: * -> *) (m :: * -> *) a b.
(a -> m b) -> t a -> m (t b)
mapM Expr -> M Expr
go [Expr]
es ReaderT RO (SupplyT Id) (Type -> Expr)
-> ReaderT RO (SupplyT Id) Type -> M Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Type -> ReaderT RO (SupplyT Id) Type
forall (m :: * -> *) a. Monad m => a -> m a
return Type
t
ETuple [Expr]
es       -> [Expr] -> Expr
ETuple  ([Expr] -> Expr) -> ReaderT RO (SupplyT Id) [Expr] -> M Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<\$> (Expr -> M Expr) -> [Expr] -> ReaderT RO (SupplyT Id) [Expr]
forall (t :: * -> *) (m :: * -> *) a b.
(a -> m b) -> t a -> m (t b)
mapM Expr -> M Expr
go [Expr]
es
ERec RecordMap Ident Expr
fs         -> RecordMap Ident Expr -> Expr
ERec    (RecordMap Ident Expr -> Expr)
-> ReaderT RO (SupplyT Id) (RecordMap Ident Expr) -> M Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<\$> (Expr -> M Expr)
-> RecordMap Ident Expr
-> ReaderT RO (SupplyT Id) (RecordMap Ident Expr)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse Expr -> M Expr
go RecordMap Ident Expr
fs
ESel Expr
e Selector
s        -> Expr -> Selector -> Expr
ESel    (Expr -> Selector -> Expr)
-> M Expr -> ReaderT RO (SupplyT Id) (Selector -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<\$> Expr -> M Expr
go Expr
e  ReaderT RO (SupplyT Id) (Selector -> Expr)
-> ReaderT RO (SupplyT Id) Selector -> M Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Selector -> ReaderT RO (SupplyT Id) Selector
forall (m :: * -> *) a. Monad m => a -> m a
return Selector
s
ESet Type
ty Expr
e Selector
s Expr
v   -> Type -> Expr -> Selector -> Expr -> Expr
ESet Type
ty (Expr -> Selector -> Expr -> Expr)
-> M Expr -> ReaderT RO (SupplyT Id) (Selector -> Expr -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<\$> Expr -> M Expr
go Expr
e  ReaderT RO (SupplyT Id) (Selector -> Expr -> Expr)
-> ReaderT RO (SupplyT Id) Selector
-> ReaderT RO (SupplyT Id) (Expr -> Expr)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Selector -> ReaderT RO (SupplyT Id) Selector
forall (m :: * -> *) a. Monad m => a -> m a
return Selector
s ReaderT RO (SupplyT Id) (Expr -> Expr) -> M Expr -> M Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expr -> M Expr
go Expr
v
EIf Expr
e1 Expr
e2 Expr
e3    -> Expr -> Expr -> Expr -> Expr
EIf     (Expr -> Expr -> Expr -> Expr)
-> M Expr -> ReaderT RO (SupplyT Id) (Expr -> Expr -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<\$> Expr -> M Expr
go Expr
e1 ReaderT RO (SupplyT Id) (Expr -> Expr -> Expr)
-> M Expr -> ReaderT RO (SupplyT Id) (Expr -> Expr)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expr -> M Expr
go Expr
e2 ReaderT RO (SupplyT Id) (Expr -> Expr) -> M Expr -> M Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expr -> M Expr
go Expr
e3

EComp Type
len Type
t Expr
e [[Match]]
mss -> Type -> Type -> Expr -> [[Match]] -> Expr
EComp Type
len Type
t (Expr -> [[Match]] -> Expr)
-> M Expr -> ReaderT RO (SupplyT Id) ([[Match]] -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<\$> Expr -> M Expr
go Expr
e  ReaderT RO (SupplyT Id) ([[Match]] -> Expr)
-> ReaderT RO (SupplyT Id) [[Match]] -> M Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ([Match] -> ReaderT RO (SupplyT Id) [Match])
-> [[Match]] -> ReaderT RO (SupplyT Id) [[Match]]
forall (t :: * -> *) (m :: * -> *) a b.
(a -> m b) -> t a -> m (t b)
mapM ((Match -> ReaderT RO (SupplyT Id) Match)
-> [Match] -> ReaderT RO (SupplyT Id) [Match]
forall (t :: * -> *) (m :: * -> *) a b.
(a -> m b) -> t a -> m (t b)
mapM (RewMap -> Match -> ReaderT RO (SupplyT Id) Match
rewM RewMap
rews)) [[Match]]
mss
EVar Name
_          -> Expr -> M Expr
forall (m :: * -> *) a. Monad m => a -> m a
return Expr
expr

ETAbs TParam
x Expr
e       -> TParam -> Expr -> Expr
ETAbs TParam
x  (Expr -> Expr) -> M Expr -> M Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<\$> Expr -> M Expr
go Expr
e

EApp Expr
e1 Expr
e2      -> Expr -> Expr -> Expr
EApp     (Expr -> Expr -> Expr)
-> M Expr -> ReaderT RO (SupplyT Id) (Expr -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<\$> Expr -> M Expr
go Expr
e1 ReaderT RO (SupplyT Id) (Expr -> Expr) -> M Expr -> M Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expr -> M Expr
go Expr
e2
EAbs Name
x Type
t Expr
e      -> Name -> Type -> Expr -> Expr
EAbs Name
x Type
t (Expr -> Expr) -> M Expr -> M Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<\$> Expr -> M Expr
go Expr
e

EProofAbs Type
x Expr
e   -> Type -> Expr -> Expr
EProofAbs Type
x (Expr -> Expr) -> M Expr -> M Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<\$> Expr -> M Expr
go Expr
e

EWhere Expr
e [DeclGroup]
dgs    -> Expr -> [DeclGroup] -> Expr
EWhere      (Expr -> [DeclGroup] -> Expr)
-> M Expr -> ReaderT RO (SupplyT Id) ([DeclGroup] -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<\$> Expr -> M Expr
go Expr
e ReaderT RO (SupplyT Id) ([DeclGroup] -> Expr)
-> ReaderT RO (SupplyT Id) [DeclGroup] -> M Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ReaderT RO (SupplyT Id) [DeclGroup]
-> ReaderT RO (SupplyT Id) [DeclGroup]
forall a. M a -> M a
inLocal
((DeclGroup -> ReaderT RO (SupplyT Id) DeclGroup)
-> [DeclGroup] -> ReaderT RO (SupplyT Id) [DeclGroup]
forall (t :: * -> *) (m :: * -> *) a b.
(a -> m b) -> t a -> m (t b)
mapM (RewMap -> DeclGroup -> ReaderT RO (SupplyT Id) DeclGroup
rewDeclGroup RewMap
rews) [DeclGroup]
dgs)

rewM :: RewMap -> Match -> M Match
rewM :: RewMap -> Match -> ReaderT RO (SupplyT Id) Match
rewM RewMap
rews Match
ma =
case Match
ma of
From Name
x Type
len Type
t Expr
e -> Name -> Type -> Type -> Expr -> Match
From Name
x Type
len Type
t (Expr -> Match) -> M Expr -> ReaderT RO (SupplyT Id) Match
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<\$> RewMap -> Expr -> M Expr
rewE RewMap
rews Expr
e

-- These are not recursive.
Let Decl
d      -> Decl -> Match
Let (Decl -> Match)
-> ReaderT RO (SupplyT Id) Decl -> ReaderT RO (SupplyT Id) Match
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<\$> RewMap -> Decl -> ReaderT RO (SupplyT Id) Decl
rewD RewMap
rews Decl
d

rewD :: RewMap -> Decl -> M Decl
rewD :: RewMap -> Decl -> ReaderT RO (SupplyT Id) Decl
rewD RewMap
rews Decl
d = do DeclDef
e <- RewMap -> DeclDef -> M DeclDef
rewDef RewMap
rews (Decl -> DeclDef
dDefinition Decl
d)
Decl -> ReaderT RO (SupplyT Id) Decl
forall (m :: * -> *) a. Monad m => a -> m a
return Decl
d { dDefinition :: DeclDef
dDefinition = DeclDef
e }

rewDef :: RewMap -> DeclDef -> M DeclDef
rewDef :: RewMap -> DeclDef -> M DeclDef
rewDef RewMap
rews (DExpr Expr
e) = Expr -> DeclDef
DExpr (Expr -> DeclDef) -> M Expr -> M DeclDef
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<\$> RewMap -> Expr -> M Expr
rewE RewMap
rews Expr
e
rewDef RewMap
_    DeclDef
DPrim     = DeclDef -> M DeclDef
forall (m :: * -> *) a. Monad m => a -> m a
return DeclDef
DPrim

rewDeclGroup :: RewMap -> DeclGroup -> M DeclGroup
rewDeclGroup :: RewMap -> DeclGroup -> ReaderT RO (SupplyT Id) DeclGroup
rewDeclGroup RewMap
rews DeclGroup
dg =
case DeclGroup
dg of
NonRecursive Decl
d -> Decl -> DeclGroup
NonRecursive (Decl -> DeclGroup)
-> ReaderT RO (SupplyT Id) Decl
-> ReaderT RO (SupplyT Id) DeclGroup
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<\$> RewMap -> Decl -> ReaderT RO (SupplyT Id) Decl
rewD RewMap
rews Decl
d
Recursive [Decl]
ds ->
do let ([Decl]
leave,[(Decl, [TParam], [Type], Expr)]
rew) = [Either Decl (Decl, [TParam], [Type], Expr)]
-> ([Decl], [(Decl, [TParam], [Type], Expr)])
forall a b. [Either a b] -> ([a], [b])
partitionEithers ((Decl -> Either Decl (Decl, [TParam], [Type], Expr))
-> [Decl] -> [Either Decl (Decl, [TParam], [Type], Expr)]
forall a b. (a -> b) -> [a] -> [b]
map Decl -> Either Decl (Decl, [TParam], [Type], Expr)
consider [Decl]
ds)
rewGroups :: [[(Decl, [TParam], [Type], Expr)]]
rewGroups   = ((Decl, [TParam], [Type], Expr)
-> (Decl, [TParam], [Type], Expr) -> Bool)
-> [(Decl, [TParam], [Type], Expr)]
-> [[(Decl, [TParam], [Type], Expr)]]
forall a. (a -> a -> Bool) -> [a] -> [[a]]
groupBy (Decl, [TParam], [Type], Expr)
-> (Decl, [TParam], [Type], Expr) -> Bool
forall a a a d a d.
(Eq a, Eq a) =>
(a, a, a, d) -> (a, a, a, d) -> Bool
sameTParams
([(Decl, [TParam], [Type], Expr)]
-> [[(Decl, [TParam], [Type], Expr)]])
-> [(Decl, [TParam], [Type], Expr)]
-> [[(Decl, [TParam], [Type], Expr)]]
forall a b. (a -> b) -> a -> b
\$ ((Decl, [TParam], [Type], Expr)
-> (Decl, [TParam], [Type], Expr) -> Ordering)
-> [(Decl, [TParam], [Type], Expr)]
-> [(Decl, [TParam], [Type], Expr)]
forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy (Decl, [TParam], [Type], Expr)
-> (Decl, [TParam], [Type], Expr) -> Ordering
forall a b a d a d.
(Ord a, Ord b) =>
(a, b, a, d) -> (a, b, a, d) -> Ordering
compareTParams [(Decl, [TParam], [Type], Expr)]
rew
[Decl]
ds1 <- (Decl -> ReaderT RO (SupplyT Id) Decl)
-> [Decl] -> ReaderT RO (SupplyT Id) [Decl]
forall (t :: * -> *) (m :: * -> *) a b.
(a -> m b) -> t a -> m (t b)
mapM (RewMap -> Decl -> ReaderT RO (SupplyT Id) Decl
rewD RewMap
rews) [Decl]
leave
[[Decl]]
ds2 <- ([(Decl, [TParam], [Type], Expr)]
-> ReaderT RO (SupplyT Id) [Decl])
-> [[(Decl, [TParam], [Type], Expr)]]
-> ReaderT RO (SupplyT Id) [[Decl]]
forall (t :: * -> *) (m :: * -> *) a b.
(a -> m b) -> t a -> m (t b)
mapM [(Decl, [TParam], [Type], Expr)] -> ReaderT RO (SupplyT Id) [Decl]
rewSame [[(Decl, [TParam], [Type], Expr)]]
rewGroups
DeclGroup -> ReaderT RO (SupplyT Id) DeclGroup
forall (m :: * -> *) a. Monad m => a -> m a
return (DeclGroup -> ReaderT RO (SupplyT Id) DeclGroup)
-> DeclGroup -> ReaderT RO (SupplyT Id) DeclGroup
forall a b. (a -> b) -> a -> b
\$ [Decl] -> DeclGroup
Recursive ([Decl]
ds1 [Decl] -> [Decl] -> [Decl]
forall a. [a] -> [a] -> [a]
++ [[Decl]] -> [Decl]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Decl]]
ds2)

where
sameTParams :: (a, a, a, d) -> (a, a, a, d) -> Bool
sameTParams    (a
_,a
tps1,a
x,d
_) (a
_,a
tps2,a
y,d
_) = a
tps1 a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
tps2 Bool -> Bool -> Bool
&& a
x a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
y
compareTParams :: (a, b, a, d) -> (a, b, a, d) -> Ordering
compareTParams (a
_,b
tps1,a
x,d
_) (a
_,b
tps2,a
y,d
_) = (a, b) -> (a, b) -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (a
x,b
tps1) (a
y,b
tps2)

consider :: Decl -> Either Decl (Decl, [TParam], [Type], Expr)
consider Decl
d   =
case Decl -> DeclDef
dDefinition Decl
d of
DeclDef
DPrim   -> Decl -> Either Decl (Decl, [TParam], [Type], Expr)
forall a b. a -> Either a b
Left Decl
d
DExpr Expr
e -> let ([TParam]
tps,[Type]
props,Expr
e') = Expr -> ([TParam], [Type], Expr)
splitTParams Expr
e
in if Bool -> Bool
not ([TParam] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [TParam]
tps) Bool -> Bool -> Bool
&& Expr -> Bool
notFun Expr
e'
then (Decl, [TParam], [Type], Expr)
-> Either Decl (Decl, [TParam], [Type], Expr)
forall a b. b -> Either a b
Right (Decl
d, [TParam]
tps, [Type]
props, Expr
e')
else Decl -> Either Decl (Decl, [TParam], [Type], Expr)
forall a b. a -> Either a b
Left Decl
d

rewSame :: [(Decl, [TParam], [Type], Expr)] -> ReaderT RO (SupplyT Id) [Decl]
rewSame [(Decl, [TParam], [Type], Expr)]
ds =
do [(Decl, Name, Expr)]
new <- [(Decl, [TParam], [Type], Expr)]
-> ((Decl, [TParam], [Type], Expr)
-> ReaderT RO (SupplyT Id) (Decl, Name, Expr))
-> ReaderT RO (SupplyT Id) [(Decl, Name, Expr)]
forall (t :: * -> *) (m :: * -> *) a b.
t a -> (a -> m b) -> m (t b)
forM [(Decl, [TParam], [Type], Expr)]
ds (((Decl, [TParam], [Type], Expr)
-> ReaderT RO (SupplyT Id) (Decl, Name, Expr))
-> ReaderT RO (SupplyT Id) [(Decl, Name, Expr)])
-> ((Decl, [TParam], [Type], Expr)
-> ReaderT RO (SupplyT Id) (Decl, Name, Expr))
-> ReaderT RO (SupplyT Id) [(Decl, Name, Expr)]
forall a b. (a -> b) -> a -> b
\$ \(Decl
d,[TParam]
_,[Type]
_,Expr
e) ->
do Name
x <- M Name
newName
(Decl, Name, Expr) -> ReaderT RO (SupplyT Id) (Decl, Name, Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Decl
d, Name
x, Expr
e)
let (Decl
_,[TParam]
tps,[Type]
props,Expr
_) : [(Decl, [TParam], [Type], Expr)]
_ = [(Decl, [TParam], [Type], Expr)]
ds
tys :: [Type]
tys            = (TParam -> Type) -> [TParam] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (TVar -> Type
TVar (TVar -> Type) -> (TParam -> TVar) -> TParam -> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TParam -> TVar
tpVar) [TParam]
tps
proofNum :: Int
proofNum       = [Type] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
props
addRew :: (Decl, a, c) -> m a -> m a
d,a
x,c
_) = (Name, [Type], Int) -> a -> m a -> m a
forall (m :: * -> *) k a. TrieMap m k => k -> a -> m a -> m a
insertTM (Decl -> Name
dName Decl
d,[Type]
tys,Int
proofNum) a
x
newRews :: RewMap
newRews        = ((Decl, Name, Expr) -> RewMap -> RewMap)
-> RewMap -> [(Decl, Name, Expr)] -> RewMap
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (Decl, Name, Expr) -> RewMap -> RewMap
forall (m :: * -> *) a c.
TrieMap m (Name, [Type], Int) =>
(Decl, a, c) -> m a -> m a
rews [(Decl, Name, Expr)]
new

[(Decl, Decl)]
newDs <- [(Decl, Name, Expr)]
-> ((Decl, Name, Expr) -> ReaderT RO (SupplyT Id) (Decl, Decl))
-> ReaderT RO (SupplyT Id) [(Decl, Decl)]
forall (t :: * -> *) (m :: * -> *) a b.
t a -> (a -> m b) -> m (t b)
forM [(Decl, Name, Expr)]
new (((Decl, Name, Expr) -> ReaderT RO (SupplyT Id) (Decl, Decl))
-> ReaderT RO (SupplyT Id) [(Decl, Decl)])
-> ((Decl, Name, Expr) -> ReaderT RO (SupplyT Id) (Decl, Decl))
-> ReaderT RO (SupplyT Id) [(Decl, Decl)]
forall a b. (a -> b) -> a -> b
\$ \(Decl
d,Name
newN,Expr
e) ->
do Expr
e1 <- RewMap -> Expr -> M Expr
rewE RewMap
newRews Expr
e
(Decl, Decl) -> ReaderT RO (SupplyT Id) (Decl, Decl)
forall (m :: * -> *) a. Monad m => a -> m a
return ( Decl
d
, Decl
d { dName :: Name
dName        = Name
newN
, dSignature :: Schema
dSignature   = (Decl -> Schema
dSignature Decl
d)
{ sVars :: [TParam]
sVars = [], sProps :: [Type]
sProps = [] }
, dDefinition :: DeclDef
dDefinition  = Expr -> DeclDef
DExpr Expr
e1
}
)

case [(Decl, Decl)]
newDs of
[(Decl
f,Decl
f')] ->
[Decl] -> ReaderT RO (SupplyT Id) [Decl]
forall (m :: * -> *) a. Monad m => a -> m a
return  [ Decl
f { dDefinition :: DeclDef
dDefinition =
let newBody :: Expr
newBody = Name -> Expr
EVar (Decl -> Name
dName Decl
f')
newE :: Expr
newE = Expr -> [DeclGroup] -> Expr
EWhere Expr
newBody
[ [Decl] -> DeclGroup
Recursive [Decl
f'] ]
in Expr -> DeclDef
DExpr (Expr -> DeclDef) -> Expr -> DeclDef
forall a b. (a -> b) -> a -> b
\$ (TParam -> Expr -> Expr) -> Expr -> [TParam] -> Expr
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr TParam -> Expr -> Expr
ETAbs
((Type -> Expr -> Expr) -> Expr -> [Type] -> Expr
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Type -> Expr -> Expr
EProofAbs Expr
newE [Type]
props) [TParam]
tps
}
]

[(Decl, Decl)]
_ -> do Name
tupName <- M Name
newTopOrLocalName
let ([Decl]
polyDs,[Decl]
monoDs) = [(Decl, Decl)] -> ([Decl], [Decl])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Decl, Decl)]
newDs

tupAr :: Int
tupAr  = [Decl] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Decl]
monoDs
addTPs = (Expr -> [TParam] -> Expr) -> [TParam] -> Expr -> Expr
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((TParam -> Expr -> Expr) -> Expr -> [TParam] -> Expr
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr TParam -> Expr -> Expr
ETAbs)     [TParam]
tps
(Expr -> Expr) -> (Expr -> Expr) -> Expr -> Expr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Expr -> [Type] -> Expr) -> [Type] -> Expr -> Expr
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Type -> Expr -> Expr) -> Expr -> [Type] -> Expr
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Type -> Expr -> Expr
EProofAbs) [Type]
props

-- tuple = \{a} p -> (f',g')
--                where f' = ...
--                      g' = ...
tupD :: Decl
tupD = Decl :: Name
-> Schema
-> DeclDef
-> [Pragma]
-> Bool
-> Maybe Fixity
-> Maybe Text
-> Decl
Decl
{ dName :: Name
dName       = Name
tupName
, dSignature :: Schema
dSignature  =
[TParam] -> [Type] -> Type -> Schema
Forall [TParam]
tps [Type]
props (Type -> Schema) -> Type -> Schema
forall a b. (a -> b) -> a -> b
\$
TCon -> [Type] -> Type
TCon (TC -> TCon
TC (Int -> TC
TCTuple Int
tupAr))
((Decl -> Type) -> [Decl] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (Schema -> Type
sType (Schema -> Type) -> (Decl -> Schema) -> Decl -> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Decl -> Schema
dSignature) [Decl]
monoDs)

, dDefinition :: DeclDef
dDefinition =
Expr -> DeclDef
DExpr  (Expr -> DeclDef) -> Expr -> DeclDef
forall a b. (a -> b) -> a -> b
\$
Expr -> Expr
addTPs (Expr -> Expr) -> Expr -> Expr
forall a b. (a -> b) -> a -> b
\$
Expr -> [DeclGroup] -> Expr
EWhere ([Expr] -> Expr
ETuple [ Name -> Expr
EVar (Decl -> Name
dName Decl
d) | Decl
d <- [Decl]
monoDs ])
[ [Decl] -> DeclGroup
Recursive [Decl]
monoDs ]

, dPragmas :: [Pragma]
dPragmas    = [] -- ?

, dInfix :: Bool
dInfix = Bool
False
, dFixity :: Maybe Fixity
dFixity = Maybe Fixity
forall a. Maybe a
Nothing
, dDoc :: Maybe Text
dDoc = Maybe Text
forall a. Maybe a
Nothing
}

mkProof :: Expr -> p -> Expr
mkProof Expr
e p
_ = Expr -> Expr
EProofApp Expr
e

-- f = \{a} (p) -> (tuple @a p). n

mkFunDef :: Int -> Decl -> Decl
mkFunDef Int
n Decl
f =
Decl
f { dDefinition :: DeclDef
dDefinition =
Expr -> DeclDef
DExpr  (Expr -> DeclDef) -> Expr -> DeclDef
forall a b. (a -> b) -> a -> b
\$
Expr -> Expr
addTPs (Expr -> Expr) -> Expr -> Expr
forall a b. (a -> b) -> a -> b
\$ Expr -> Selector -> Expr
ESel ( (Expr -> [Type] -> Expr) -> [Type] -> Expr -> Expr
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Expr -> Type -> Expr) -> Expr -> [Type] -> Expr
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Expr -> Type -> Expr
forall p. Expr -> p -> Expr
mkProof) [Type]
props
(Expr -> Expr) -> Expr -> Expr
forall a b. (a -> b) -> a -> b
\$ (Expr -> [Type] -> Expr) -> [Type] -> Expr -> Expr
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Expr -> Type -> Expr) -> Expr -> [Type] -> Expr
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Expr -> Type -> Expr
ETApp) [Type]
tys
(Expr -> Expr) -> Expr -> Expr
forall a b. (a -> b) -> a -> b
\$ Name -> Expr
EVar Name
tupName
) (Int -> Maybe Int -> Selector
TupleSel Int
n (Int -> Maybe Int
forall a. a -> Maybe a
Just Int
tupAr))
}

[Decl] -> ReaderT RO (SupplyT Id) [Decl]
forall (m :: * -> *) a. Monad m => a -> m a
return (Decl
tupD Decl -> [Decl] -> [Decl]
forall a. a -> [a] -> [a]
: (Int -> Decl -> Decl) -> [Int] -> [Decl] -> [Decl]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> Decl -> Decl
mkFunDef [ Int
0 .. ] [Decl]
polyDs)

--------------------------------------------------------------------------------
splitTParams :: Expr -> ([TParam], [Prop], Expr)
splitTParams :: Expr -> ([TParam], [Type], Expr)
splitTParams Expr
e = let ([TParam]
tps, Expr
e1)   = (Expr -> Maybe (TParam, Expr)) -> Expr -> ([TParam], Expr)
forall a b. (a -> Maybe (b, a)) -> a -> ([b], a)
splitWhile Expr -> Maybe (TParam, Expr)
splitTAbs Expr
e
([Type]
props, Expr
e2) = (Expr -> Maybe (Type, Expr)) -> Expr -> ([Type], Expr)
forall a b. (a -> Maybe (b, a)) -> a -> ([b], a)
splitWhile Expr -> Maybe (Type, Expr)
splitProofAbs Expr
e1
in ([TParam]
tps,[Type]
props,Expr
e2)

-- returns type instantitaion and how many "proofs" were there
splitTApp :: Expr -> Int -> (Expr, [Type], Int)
splitTApp :: Expr -> Int -> (Expr, [Type], Int)
splitTApp (EProofApp Expr
e) Int
n = Expr -> Int -> (Expr, [Type], Int)
splitTApp Expr
e (Int -> (Expr, [Type], Int)) -> Int -> (Expr, [Type], Int)
forall a b. (a -> b) -> a -> b
\$! (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
splitTApp Expr
e0 Int
n            = let (Expr
e1,[Type]
ts) = Expr -> [Type] -> (Expr, [Type])
splitTy Expr
e0 []
in (Expr
e1, [Type]
ts, Int
n)
where
splitTy :: Expr -> [Type] -> (Expr, [Type])
splitTy (ETApp Expr
e Type
t) [Type]
ts = Expr -> [Type] -> (Expr, [Type])
splitTy Expr
e (Type
tType -> [Type] -> [Type]
forall a. a -> [a] -> [a]
:[Type]
ts)
splitTy Expr
e [Type]
ts           = (Expr
e,[Type]
ts)

notFun :: Expr -> Bool
notFun :: Expr -> Bool
notFun (EAbs {})       = Bool
False
notFun (EProofAbs Type
_ Expr
e) = Expr -> Bool
notFun Expr
e
notFun Expr
_               = Bool
True

```