-- |
-- Module      :  Cryptol.Transform.MonoValues
-- Copyright   :  (c) 2013-2016 Galois, Inc.
-- License     :  BSD3
-- 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 #-}
{-# LANGUAGE OverloadedStrings #-}
module Cryptol.Transform.MonoValues (rewModule) where

import Cryptol.ModuleSystem.Name
        (SupplyT,liftSupply,Supply,mkDeclared,NameSource(..),ModPath(..))
import Cryptol.Parser.Position (emptyRange)
import Cryptol.TypeCheck.AST hiding (splitTApp) -- XXX: just use this one
import Cryptol.TypeCheck.TypeMap
import Cryptol.Utils.Ident(Namespace(..))
import Data.List(sortBy)
import Data.Either(partitionEithers)
import Data.Map (Map)
import qualified Data.List.NonEmpty as NE
import MonadLib hiding (mapM)

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 :: forall a. 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 a. Map Name a
forall (m :: * -> *) k a. TrieMap m k => m a
emptyTM

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

  lookupTM :: forall a. (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 a. Name -> Map Name a -> Maybe 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 a. [Type] -> List TypeMap a -> Maybe 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 a. 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 :: forall a.
(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 a. Name -> (Maybe a -> Maybe a) -> Map Name a -> Map Name 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 a. a -> Maybe 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 a. m a
forall (m :: * -> *) k a. TrieMap m k => m a
emptyTM) m (m a)
forall a. 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 a. [Type] -> (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 [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 a. a -> Maybe 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 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 a. 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 :: forall a. (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 a. (a -> a -> a) -> Map Name a -> Map Name a -> Map Name 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 a.
(a -> a -> a) -> List TypeMap a -> List TypeMap a -> List TypeMap 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 a. (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 :: forall a. 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 a. Map Name a -> [(Name, 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 a. List TypeMap a -> [([Type], 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 a. Map Int a -> [(Int, a)]
forall (m :: * -> *) k a. TrieMap m k => m a -> [(k, a)]
toListTM Map Int a
pM ]

  mapMaybeWithKeyTM :: forall a b.
((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 a b. (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 -> ModuleG ModName -> (ModuleG ModName, Supply)
rewModule Supply
s ModuleG ModName
m = ReaderT RO (SupplyT Id) (ModuleG ModName)
-> RO -> Supply -> (ModuleG ModName, Supply)
forall (m :: * -> *) a r. RunM m a r => m a -> r
runM ReaderT RO (SupplyT Id) (ModuleG ModName)
body (ModName -> RO
TopModule (ModuleG ModName -> ModName
forall mname. ModuleG mname -> mname
mName ModuleG ModName
m)) Supply
s
  where
  body :: ReaderT RO (SupplyT Id) (ModuleG ModName)
body = do [DeclGroup]
ds <- (DeclGroup -> ReaderT RO (SupplyT Id) DeclGroup)
-> [DeclGroup] -> ReaderT RO (SupplyT Id) [DeclGroup]
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 (RewMap -> DeclGroup -> ReaderT RO (SupplyT Id) DeclGroup
rewDeclGroup RewMap
forall a. RewMap' a
forall (m :: * -> *) k a. TrieMap m k => m a
emptyTM) (ModuleG ModName -> [DeclGroup]
forall mname. ModuleG mname -> [DeclGroup]
mDecls ModuleG ModName
m)
            ModuleG ModName -> ReaderT RO (SupplyT Id) (ModuleG ModName)
forall a. a -> ReaderT RO (SupplyT Id) a
forall (m :: * -> *) a. Monad m => a -> m a
return ModuleG ModName
m { mDecls = ds }

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

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

-- | 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
ask
     (Supply -> (Name, Supply)) -> M Name
forall a. (Supply -> (a, Supply)) -> ReaderT RO (SupplyT Id) a
forall (m :: * -> *) a. FreshM m => (Supply -> (a, Supply)) -> m a
liftSupply (Namespace
-> RO
-> NameSource
-> Ident
-> Maybe Fixity
-> Range
-> Supply
-> (Name, Supply)
mkDeclared Namespace
NSValue 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 :: forall a. 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 a. (Name, [Type], Int) -> RewMap' a -> Maybe a
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 a. a -> Maybe a
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 a b.
ReaderT RO (SupplyT Id) (a -> b)
-> ReaderT RO (SupplyT Id) a -> ReaderT RO (SupplyT Id) b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Type -> ReaderT RO (SupplyT Id) Type
forall a. a -> ReaderT RO (SupplyT Id) a
forall (m :: * -> *) a. Monad m => a -> m a
return Type
t
                          Just Expr
yes -> Expr -> M Expr
forall a. a -> ReaderT RO (SupplyT Id) a
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 a. a -> ReaderT RO (SupplyT Id) a
forall (m :: * -> *) a. Monad m => a -> m a
return Expr
yes

      ELocated Range
r Expr
t    -> Range -> Expr -> Expr
ELocated Range
r (Expr -> Expr) -> M Expr -> M Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr -> M Expr
go Expr
t
      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.
(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 Expr -> M Expr
go [Expr]
es ReaderT RO (SupplyT Id) (Type -> Expr)
-> ReaderT RO (SupplyT Id) Type -> M Expr
forall a b.
ReaderT RO (SupplyT Id) (a -> b)
-> ReaderT RO (SupplyT Id) a -> ReaderT RO (SupplyT Id) b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Type -> ReaderT RO (SupplyT Id) Type
forall a. a -> ReaderT RO (SupplyT Id) a
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.
(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 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)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> RecordMap Ident a -> f (RecordMap Ident 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 a b.
ReaderT RO (SupplyT Id) (a -> b)
-> ReaderT RO (SupplyT Id) a -> ReaderT RO (SupplyT Id) b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Selector -> ReaderT RO (SupplyT Id) Selector
forall a. a -> ReaderT RO (SupplyT Id) a
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 a b.
ReaderT RO (SupplyT Id) (a -> b)
-> ReaderT RO (SupplyT Id) a -> ReaderT RO (SupplyT Id) b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Selector -> ReaderT RO (SupplyT Id) Selector
forall a. a -> ReaderT RO (SupplyT Id) a
forall (m :: * -> *) a. Monad m => a -> m a
return Selector
s ReaderT RO (SupplyT Id) (Expr -> Expr) -> M Expr -> M Expr
forall a b.
ReaderT RO (SupplyT Id) (a -> b)
-> ReaderT RO (SupplyT Id) a -> ReaderT RO (SupplyT Id) b
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 a b.
ReaderT RO (SupplyT Id) (a -> b)
-> ReaderT RO (SupplyT Id) a -> ReaderT RO (SupplyT Id) b
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 a b.
ReaderT RO (SupplyT Id) (a -> b)
-> ReaderT RO (SupplyT Id) a -> ReaderT RO (SupplyT Id) b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expr -> M Expr
go Expr
e3
      ECase Expr
e Map Ident CaseAlt
as Maybe CaseAlt
d    -> Expr -> Map Ident CaseAlt -> Maybe CaseAlt -> Expr
ECase   (Expr -> Map Ident CaseAlt -> Maybe CaseAlt -> Expr)
-> M Expr
-> ReaderT
     RO (SupplyT Id) (Map Ident CaseAlt -> Maybe CaseAlt -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr -> M Expr
go Expr
e  ReaderT
  RO (SupplyT Id) (Map Ident CaseAlt -> Maybe CaseAlt -> Expr)
-> ReaderT RO (SupplyT Id) (Map Ident CaseAlt)
-> ReaderT RO (SupplyT Id) (Maybe CaseAlt -> Expr)
forall a b.
ReaderT RO (SupplyT Id) (a -> b)
-> ReaderT RO (SupplyT Id) a -> ReaderT RO (SupplyT Id) b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (CaseAlt -> ReaderT RO (SupplyT Id) CaseAlt)
-> Map Ident CaseAlt -> ReaderT RO (SupplyT Id) (Map Ident CaseAlt)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Map Ident a -> f (Map Ident b)
traverse (RewMap -> CaseAlt -> ReaderT RO (SupplyT Id) CaseAlt
rewCase RewMap
rews) Map Ident CaseAlt
as
                                           ReaderT RO (SupplyT Id) (Maybe CaseAlt -> Expr)
-> ReaderT RO (SupplyT Id) (Maybe CaseAlt) -> M Expr
forall a b.
ReaderT RO (SupplyT Id) (a -> b)
-> ReaderT RO (SupplyT Id) a -> ReaderT RO (SupplyT Id) b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (CaseAlt -> ReaderT RO (SupplyT Id) CaseAlt)
-> Maybe CaseAlt -> ReaderT RO (SupplyT Id) (Maybe CaseAlt)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Maybe a -> f (Maybe b)
traverse (RewMap -> CaseAlt -> ReaderT RO (SupplyT Id) CaseAlt
rewCase RewMap
rews) Maybe CaseAlt
d

      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 a b.
ReaderT RO (SupplyT Id) (a -> b)
-> ReaderT RO (SupplyT Id) a -> ReaderT RO (SupplyT Id) b
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.
(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 ((Match -> ReaderT RO (SupplyT Id) Match)
-> [Match] -> ReaderT RO (SupplyT Id) [Match]
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 (RewMap -> Match -> ReaderT RO (SupplyT Id) Match
rewM RewMap
rews)) [[Match]]
mss
      EVar Name
_          -> Expr -> M Expr
forall a. a -> ReaderT RO (SupplyT Id) a
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 a b.
ReaderT RO (SupplyT Id) (a -> b)
-> ReaderT RO (SupplyT Id) a -> ReaderT RO (SupplyT Id) b
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 a b.
ReaderT RO (SupplyT Id) (a -> b)
-> ReaderT RO (SupplyT Id) a -> ReaderT RO (SupplyT Id) b
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.
(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 (RewMap -> DeclGroup -> ReaderT RO (SupplyT Id) DeclGroup
rewDeclGroup RewMap
rews) [DeclGroup]
dgs)

      EPropGuards [([Type], Expr)]
guards Type
ty -> [([Type], Expr)] -> Type -> Expr
EPropGuards ([([Type], Expr)] -> Type -> Expr)
-> ReaderT RO (SupplyT Id) [([Type], Expr)]
-> ReaderT RO (SupplyT Id) (Type -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (\([Type]
props, Expr
e) -> (,) ([Type] -> Expr -> ([Type], Expr))
-> ReaderT RO (SupplyT Id) [Type]
-> ReaderT RO (SupplyT Id) (Expr -> ([Type], Expr))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Type] -> ReaderT RO (SupplyT Id) [Type]
forall a. a -> ReaderT RO (SupplyT Id) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Type]
props ReaderT RO (SupplyT Id) (Expr -> ([Type], Expr))
-> M Expr -> ReaderT RO (SupplyT Id) ([Type], Expr)
forall a b.
ReaderT RO (SupplyT Id) (a -> b)
-> ReaderT RO (SupplyT Id) a -> ReaderT RO (SupplyT Id) b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expr -> M Expr
go Expr
e) (([Type], Expr) -> ReaderT RO (SupplyT Id) ([Type], Expr))
-> [([Type], Expr)] -> ReaderT RO (SupplyT Id) [([Type], Expr)]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
`traverse` [([Type], Expr)]
guards ReaderT RO (SupplyT Id) (Type -> Expr)
-> ReaderT RO (SupplyT Id) Type -> M Expr
forall a b.
ReaderT RO (SupplyT Id) (a -> b)
-> ReaderT RO (SupplyT Id) a -> ReaderT RO (SupplyT Id) b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Type -> ReaderT RO (SupplyT Id) Type
forall a. a -> ReaderT RO (SupplyT Id) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Type
ty


rewCase :: RewMap -> CaseAlt -> M CaseAlt
rewCase :: RewMap -> CaseAlt -> ReaderT RO (SupplyT Id) CaseAlt
rewCase RewMap
rew (CaseAlt [(Name, Type)]
xs Expr
e) = [(Name, Type)] -> Expr -> CaseAlt
CaseAlt [(Name, Type)]
xs (Expr -> CaseAlt) -> M Expr -> ReaderT RO (SupplyT Id) CaseAlt
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> RewMap -> Expr -> M Expr
rewE RewMap
rew Expr
e

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 a. a -> ReaderT RO (SupplyT Id) a
forall (m :: * -> *) a. Monad m => a -> m a
return Decl
d { dDefinition = 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 a. a -> ReaderT RO (SupplyT Id) a
forall (m :: * -> *) a. Monad m => a -> m a
return DeclDef
DPrim
rewDef RewMap
rews (DForeign FFIFunType
t Maybe Expr
me) = FFIFunType -> Maybe Expr -> DeclDef
DForeign FFIFunType
t (Maybe Expr -> DeclDef)
-> ReaderT RO (SupplyT Id) (Maybe Expr) -> M DeclDef
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Expr -> M Expr)
-> Maybe Expr -> ReaderT RO (SupplyT Id) (Maybe Expr)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Maybe a -> f (Maybe b)
traverse (RewMap -> Expr -> M Expr
rewE RewMap
rews) Maybe Expr
me

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 :: [NonEmpty (Decl, [TParam], [Type], Expr)]
rewGroups   = ((Decl, [TParam], [Type], Expr)
 -> (Decl, [TParam], [Type], Expr) -> Bool)
-> [(Decl, [TParam], [Type], Expr)]
-> [NonEmpty (Decl, [TParam], [Type], Expr)]
forall (f :: * -> *) a.
Foldable f =>
(a -> a -> Bool) -> f a -> [NonEmpty a]
NE.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)]
 -> [NonEmpty (Decl, [TParam], [Type], Expr)])
-> [(Decl, [TParam], [Type], Expr)]
-> [NonEmpty (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.
(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 (RewMap -> Decl -> ReaderT RO (SupplyT Id) Decl
rewD RewMap
rews) [Decl]
leave
         [[Decl]]
ds2 <- (NonEmpty (Decl, [TParam], [Type], Expr)
 -> ReaderT RO (SupplyT Id) [Decl])
-> [NonEmpty (Decl, [TParam], [Type], Expr)]
-> ReaderT RO (SupplyT Id) [[Decl]]
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 NonEmpty (Decl, [TParam], [Type], Expr)
-> ReaderT RO (SupplyT Id) [Decl]
rewSame [NonEmpty (Decl, [TParam], [Type], Expr)]
rewGroups
         DeclGroup -> ReaderT RO (SupplyT Id) DeclGroup
forall a. a -> ReaderT RO (SupplyT Id) a
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
      DForeign FFIFunType
_ Maybe Expr
me -> case Maybe Expr
me of
                         Maybe Expr
Nothing -> Decl -> Either Decl (Decl, [TParam], [Type], Expr)
forall a b. a -> Either a b
Left Decl
d
                         Just Expr
e  -> Expr -> Either Decl (Decl, [TParam], [Type], Expr)
conExpr Expr
e
      DExpr Expr
e       -> Expr -> Either Decl (Decl, [TParam], [Type], Expr)
conExpr Expr
e
    where
    conExpr :: Expr -> Either Decl (Decl, [TParam], [Type], Expr)
conExpr Expr
e =
      let ([TParam]
tps,[Type]
props,Expr
e') = Expr -> ([TParam], [Type], Expr)
splitTParams Expr
e
      in if Bool -> Bool
not ([TParam] -> Bool
forall a. [a] -> 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 :: NonEmpty (Decl, [TParam], [Type], Expr)
-> ReaderT RO (SupplyT Id) [Decl]
rewSame NonEmpty (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.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (NonEmpty (Decl, [TParam], [Type], Expr)
-> [(Decl, [TParam], [Type], Expr)]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty (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 a. a -> ReaderT RO (SupplyT Id) a
forall (m :: * -> *) a. Monad m => a -> m a
return (Decl
d, Name
x, Expr
e)
        let (Decl
_,[TParam]
tps,[Type]
props,Expr
_) NE.:| [(Decl, [TParam], [Type], Expr)]
_ = NonEmpty (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 a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
props
            addRew :: (Decl, a, c) -> m a -> m a
addRew (Decl
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 a b. (a -> b -> b) -> b -> [a] -> b
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
addRew RewMap
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.
(Traversable t, Monad m) =>
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 a. a -> ReaderT RO (SupplyT Id) a
forall (m :: * -> *) a. Monad m => a -> m a
return ( Decl
d
                             , Decl
d { dName        = newN
                                 , dSignature   = (dSignature d)
                                         { sVars = [], sProps = [] }
                                 , dDefinition  = DExpr e1
                                 }
                             )

        case [(Decl, Decl)]
newDs of
          [(Decl
f,Decl
f')] ->
              [Decl] -> ReaderT RO (SupplyT Id) [Decl]
forall a. a -> ReaderT RO (SupplyT Id) a
forall (m :: * -> *) a. Monad m => a -> m a
return  [ Decl
f { dDefinition =
                                let newBody = Name -> Expr
EVar (Decl -> Name
dName Decl
f')
                                    newE = Expr -> [DeclGroup] -> Expr
EWhere Expr
newBody
                                              [ [Decl] -> DeclGroup
Recursive [Decl
f'] ]
                                in DExpr $ foldr ETAbs
                                   (foldr EProofAbs newE props) 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 a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Decl]
monoDs
                      addTPs :: Expr -> Expr
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 a b. (a -> b -> b) -> b -> [a] -> b
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 a b. (a -> b -> b) -> b -> [a] -> b
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
                        { 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 =
                              DExpr  $
                              addTPs $ ESel ( flip (foldl mkProof) props
                                            $ flip (foldl ETApp) tys
                                            $ EVar tupName
                                            ) (TupleSel n (Just tupAr))
                          }

                  [Decl] -> ReaderT RO (SupplyT Id) [Decl]
forall a. a -> ReaderT RO (SupplyT Id) a
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