-- |
-- 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  = forall a. Map Name (TypesMap (Map Int a)) -> RewMap' a
RM 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) = 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 <- 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 <- forall (m :: * -> *) k a. TrieMap m k => k -> m a -> Maybe a
lookupTM [Type]
ts TypesMap (Map Int a)
tM
                                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) = forall a. Map Name (TypesMap (Map Int a)) -> RewMap' a
RM (forall (m :: * -> *) k a.
TrieMap m k =>
k -> (Maybe a -> Maybe a) -> m a -> m a
alterTM Name
x 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 forall a. Maybe a
Nothing
                      forall (m :: * -> *) a. Monad m => a -> m a
return (forall (m :: * -> *) k a. TrieMap m k => k -> a -> m a -> m a
insertTM [Type]
ts (forall (m :: * -> *) k a. TrieMap m k => k -> a -> m a -> m a
insertTM Int
n a
a forall (m :: * -> *) k a. TrieMap m k => m a
emptyTM) forall (m :: * -> *) k a. TrieMap m k => m a
emptyTM)

    f1 (Just m (m a)
tM) = forall a. a -> Maybe a
Just (forall (m :: * -> *) k a.
TrieMap m k =>
k -> (Maybe a -> Maybe a) -> m a -> m a
alterTM [Type]
ts 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 forall a. Maybe a
Nothing
                      forall (m :: * -> *) a. Monad m => a -> m a
return (forall (m :: * -> *) k a. TrieMap m k => k -> a -> m a -> m a
insertTM Int
n a
a forall (m :: * -> *) k a. TrieMap m k => m a
emptyTM)

    f2 (Just m a
pM) = forall a. a -> Maybe a
Just (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) = forall a. Map Name (TypesMap (Map Int a)) -> RewMap' a
RM (forall (m :: * -> *) k a.
TrieMap m k =>
(a -> a -> a) -> m a -> m a -> m a
unionTM (forall (m :: * -> *) k a.
TrieMap m k =>
(a -> a -> a) -> m a -> m a -> m a
unionTM (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)  <- 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) <- forall (m :: * -> *) k a. TrieMap m k => m a -> [(k, a)]
toListTM TypesMap (Map Int a)
tM
                                   , (Int
n,a
y)   <- 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) =
    forall a. Map Name (TypesMap (Map Int a)) -> RewMap' a
RM (forall (m :: * -> *) k a b.
TrieMap m k =>
(k -> a -> b) -> m a -> m b
mapWithKeyTM      (\Name
qn  TypesMap (Map Int a)
tm ->
        forall (m :: * -> *) k a b.
TrieMap m k =>
(k -> a -> b) -> m a -> m b
mapWithKeyTM      (\[Type]
tys Map Int a
is ->
        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 = forall (m :: * -> *) a r. RunM m a r => m a -> r
runM ReaderT RO (SupplyT Id) (ModuleG ModName)
body (ModName -> RO
TopModule (forall mname. ModuleG mname -> mname
mName ModuleG ModName
m)) Supply
s
  where
  body :: ReaderT RO (SupplyT Id) (ModuleG ModName)
body = do [DeclGroup]
ds <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (RewMap -> DeclGroup -> M DeclGroup
rewDeclGroup forall (m :: * -> *) k a. TrieMap m k => m a
emptyTM) (forall mname. ModuleG mname -> [DeclGroup]
mDecls ModuleG ModName
m)
            forall (m :: * -> *) a. Monad m => a -> m a
return ModuleG ModName
m { mDecls :: [DeclGroup]
mDecls = [DeclGroup]
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 <- forall (m :: * -> *) i. ReaderM m i => m i
ask
     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" 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  = 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 <- forall (m :: * -> *) k a. TrieMap m k => k -> m a -> Maybe a
lookupTM (Name
x,[Type]
tps,Int
n) RewMap
rews
        forall (m :: * -> *) a. Monad m => a -> m a
return (Name -> Expr
EVar Name
y)
  tryRewrite (Expr, [Type], Int)
_ = 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 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr -> M Expr
go Expr
e forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (m :: * -> *) a. Monad m => a -> m a
return Type
t
                          Just Expr
yes -> 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 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr -> M Expr
go Expr
e
                          Just Expr
yes -> forall (m :: * -> *) a. Monad m => a -> m a
return Expr
yes

      ELocated Range
r Expr
t    -> Range -> Expr -> Expr
ELocated Range
r 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   forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Expr -> M Expr
go [Expr]
es forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (m :: * -> *) a. Monad m => a -> m a
return Type
t
      ETuple [Expr]
es       -> [Expr] -> Expr
ETuple  forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(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    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr -> M Expr
go Expr
e  forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> 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 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr -> M Expr
go Expr
e  forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (m :: * -> *) a. Monad m => a -> m a
return Selector
s 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     forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr -> M Expr
go Expr
e1 forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expr -> M Expr
go Expr
e2 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 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr -> M Expr
go Expr
e  forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (RewMap -> Match -> M Match
rewM RewMap
rews)) [[Match]]
mss
      EVar Name
_          -> forall (m :: * -> *) a. Monad m => a -> m a
return Expr
expr

      ETAbs TParam
x Expr
e       -> TParam -> Expr -> Expr
ETAbs TParam
x  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     forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr -> M Expr
go Expr
e1 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 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 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      forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr -> M Expr
go Expr
e forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a. M a -> M a
inLocal
                                                  (forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (RewMap -> DeclGroup -> M DeclGroup
rewDeclGroup RewMap
rews) [DeclGroup]
dgs)

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


rewM :: RewMap -> Match -> M Match
rewM :: RewMap -> Match -> M 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 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 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> RewMap -> Decl -> M Decl
rewD RewMap
rews Decl
d


rewD :: RewMap -> Decl -> M Decl
rewD :: RewMap -> Decl -> M Decl
rewD RewMap
rews Decl
d = do DeclDef
e <- RewMap -> DeclDef -> M DeclDef
rewDef RewMap
rews (Decl -> DeclDef
dDefinition Decl
d)
                 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 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        = forall (m :: * -> *) a. Monad m => a -> m a
return DeclDef
DPrim
rewDef RewMap
_    (DForeign FFIFunType
t) = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ FFIFunType -> DeclDef
DForeign FFIFunType
t

rewDeclGroup :: RewMap -> DeclGroup -> M DeclGroup
rewDeclGroup :: RewMap -> DeclGroup -> M DeclGroup
rewDeclGroup RewMap
rews DeclGroup
dg =
  case DeclGroup
dg of
    NonRecursive Decl
d -> Decl -> DeclGroup
NonRecursive forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> RewMap -> Decl -> M Decl
rewD RewMap
rews Decl
d
    Recursive [Decl]
ds ->
      do let ([Decl]
leave,[(Decl, [TParam], [Type], Expr)]
rew) = forall a b. [Either a b] -> ([a], [b])
partitionEithers (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   = forall (f :: * -> *) a.
Foldable f =>
(a -> a -> Bool) -> f a -> [NonEmpty a]
NE.groupBy forall {a} {a} {a} {d} {a} {d}.
(Eq a, Eq a) =>
(a, a, a, d) -> (a, a, a, d) -> Bool
sameTParams
                         forall a b. (a -> b) -> a -> b
$ forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy 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 <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (RewMap -> Decl -> M Decl
rewD RewMap
rews) [Decl]
leave
         [[Decl]]
ds2 <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM NonEmpty (Decl, [TParam], [Type], Expr)
-> ReaderT RO (SupplyT Id) [Decl]
rewSame [NonEmpty (Decl, [TParam], [Type], Expr)]
rewGroups
         forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ [Decl] -> DeclGroup
Recursive ([Decl]
ds1 forall a. [a] -> [a] -> [a]
++ 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 forall a. Eq a => a -> a -> Bool
== a
tps2 Bool -> Bool -> Bool
&& a
x 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
_) = 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      -> forall a b. a -> Either a b
Left Decl
d
      DForeign FFIFunType
_ -> 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 (forall (t :: * -> *) a. Foldable t => t a -> Bool
null [TParam]
tps) Bool -> Bool -> Bool
&& Expr -> Bool
notFun Expr
e'
                        then forall a b. b -> Either a b
Right (Decl
d, [TParam]
tps, [Type]
props, Expr
e')
                        else 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 <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (forall a. NonEmpty a -> [a]
NE.toList NonEmpty (Decl, [TParam], [Type], Expr)
ds) forall a b. (a -> b) -> a -> b
$ \(Decl
d,[TParam]
_,[Type]
_,Expr
e) ->
                 do Name
x <- M Name
newName
                    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            = forall a b. (a -> b) -> [a] -> [b]
map (TVar -> Type
TVar forall b c a. (b -> c) -> (a -> b) -> a -> c
. TParam -> TVar
tpVar) [TParam]
tps
            proofNum :: Int
proofNum       = 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
_) = 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        = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr 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 <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [(Decl, Name, Expr)]
new 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
                      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')] ->
              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 forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr TParam -> Expr -> Expr
ETAbs
                                   (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) = forall a b. [(a, b)] -> ([a], [b])
unzip [(Decl, Decl)]
newDs

                      tupAr :: Int
tupAr  = forall (t :: * -> *) a. Foldable t => t a -> Int
length [Decl]
monoDs
                      addTPs :: Expr -> Expr
addTPs = forall a b c. (a -> b -> c) -> b -> a -> c
flip (forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr TParam -> Expr -> Expr
ETAbs)     [TParam]
tps
                             forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b c. (a -> b -> c) -> b -> a -> c
flip (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 forall a b. (a -> b) -> a -> b
$
                               TCon -> [Type] -> Type
TCon (TC -> TCon
TC (Int -> TC
TCTuple Int
tupAr))
                                    (forall a b. (a -> b) -> [a] -> [b]
map (Schema -> Type
sType forall b c a. (b -> c) -> (a -> b) -> a -> c
. Decl -> Schema
dSignature) [Decl]
monoDs)

                        , dDefinition :: DeclDef
dDefinition =
                            Expr -> DeclDef
DExpr  forall a b. (a -> b) -> a -> b
$
                            Expr -> Expr
addTPs 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 = forall a. Maybe a
Nothing
                        , dDoc :: Maybe Text
dDoc = 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  forall a b. (a -> b) -> a -> b
$
                              Expr -> Expr
addTPs forall a b. (a -> b) -> a -> b
$ Expr -> Selector -> Expr
ESel ( forall a b c. (a -> b -> c) -> b -> a -> c
flip (forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl forall {p}. Expr -> p -> Expr
mkProof) [Type]
props
                                            forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> b -> a -> c
flip (forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Expr -> Type -> Expr
ETApp) [Type]
tys
                                            forall a b. (a -> b) -> a -> b
$ Name -> Expr
EVar Name
tupName
                                            ) (Int -> Maybe Int -> Selector
TupleSel Int
n (forall a. a -> Maybe a
Just Int
tupAr))
                          }

                  forall (m :: * -> *) a. Monad m => a -> m a
return (Decl
tupD forall a. a -> [a] -> [a]
: 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)   = forall a b. (a -> Maybe (b, a)) -> a -> ([b], a)
splitWhile Expr -> Maybe (TParam, Expr)
splitTAbs Expr
e
                     ([Type]
props, Expr
e2) = 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 forall a b. (a -> b) -> a -> b
$! (Int
n 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
tforall 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