content: -- |
  -- 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  = RM emptyTM
  
    nullTM (RM m) = nullTM m
  
    lookupTM (x,ts,n) (RM m) = do tM <- lookupTM x m
                                  tP <- lookupTM ts tM
                                  lookupTM n tP
  
    alterTM (x,ts,n) f (RM m) = RM (alterTM x f1 m)
      where
      f1 Nothing   = do a <- f Nothing
                        return (insertTM ts (insertTM n a emptyTM) emptyTM)
  
      f1 (Just tM) = Just (alterTM ts f2 tM)
  
      f2 Nothing   = do a <- f Nothing
                        return (insertTM n a emptyTM)
  
      f2 (Just pM) = Just (alterTM n f pM)
  
    unionTM f (RM a) (RM b) = RM (unionTM (unionTM (unionTM f)) a b)
  
    toListTM (RM m) = [ ((x,ts,n),y) | (x,tM)  <- toListTM m
                                     , (ts,pM) <- toListTM tM
                                     , (n,y)   <- toListTM pM ]
  
    mapMaybeWithKeyTM f (RM m) =
      RM (mapWithKeyTM      (\qn  tm ->
          mapWithKeyTM      (\tys is ->
          mapMaybeWithKeyTM (\i   a  -> f (qn,tys,i) a) is) tm) 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 s m = runM body (TopModule (mName m)) s
    where
    body = do ds <- mapM (rewDeclGroup emptyTM) (mDecls m)
              return m { mDecls = ds }
  
  --------------------------------------------------------------------------------
  
  type M  = ReaderT RO (SupplyT Id)
  type RO = ModPath
  
  -- | Produce a fresh top-level name.
  newName :: M Name
  newName  =
    do ns <- ask
       liftSupply (mkDeclared NSValue ns SystemName "$mono" Nothing emptyRange)
  
  newTopOrLocalName :: M Name
  newTopOrLocalName  = 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  = id
  
  
  
  --------------------------------------------------------------------------------
  rewE :: RewMap -> Expr -> M Expr   -- XXX: not IO
  rewE rews = go
  
    where
    tryRewrite (EVar x,tps,n) =
       do y <- lookupTM (x,tps,n) rews
          return (EVar y)
    tryRewrite _ = Nothing
  
    go expr =
      case expr of
  
        -- Interesting cases
        ETApp e t      -> case tryRewrite (splitTApp expr 0) of
                            Nothing  -> ETApp <$> go e <*> return t
                            Just yes -> return yes
        EProofApp e    -> case tryRewrite (splitTApp e 1) of
                            Nothing  -> EProofApp <$> go e
                            Just yes -> return yes
  
        ELocated r t    -> ELocated r <$> go t
        EList es t      -> EList   <$> mapM go es <*> return t
        ETuple es       -> ETuple  <$> mapM go es
        ERec fs         -> ERec    <$> traverse go fs
        ESel e s        -> ESel    <$> go e  <*> return s
        ESet ty e s v   -> ESet ty <$> go e  <*> return s <*> go v
        EIf e1 e2 e3    -> EIf     <$> go e1 <*> go e2 <*> go e3
        ECase e as d    -> ECase   <$> go e  <*> traverse (rewCase rews) as
                                             <*> traverse (rewCase rews) d
  
        EComp len t e mss -> EComp len t <$> go e  <*> mapM (mapM (rewM rews)) mss
        EVar _          -> return expr
  
        ETAbs x e       -> ETAbs x  <$> go e
  
        EApp e1 e2      -> EApp     <$> go e1 <*> go e2
        EAbs x t e      -> EAbs x t <$> go e
  
        EProofAbs x e   -> EProofAbs x <$> go e
  
        EWhere e dgs    -> EWhere      <$> go e <*> inLocal
                                                    (mapM (rewDeclGroup rews) dgs)
  
        EPropGuards guards ty -> EPropGuards <$> (\(props, e) -> (,) <$> pure props <*> go e) `traverse` guards <*> pure ty
  
  
  rewCase :: RewMap -> CaseAlt -> M CaseAlt
  rewCase rew (CaseAlt xs e) = CaseAlt xs <$> rewE rew e
  
  rewM :: RewMap -> Match -> M Match
  rewM rews ma =
    case ma of
      From x len t e -> From x len t <$> rewE rews e
  
      -- These are not recursive.
      Let d      -> Let <$> rewD rews d
  
  
  rewD :: RewMap -> Decl -> M Decl
  rewD rews d = do e <- rewDef rews (dDefinition d)
                   return d { dDefinition = e }
  
  rewDef :: RewMap -> DeclDef -> M DeclDef
  rewDef rews (DExpr e)       = DExpr <$> rewE rews e
  rewDef _    DPrim           = return DPrim
  rewDef rews (DForeign t me) = DForeign t <$> traverse (rewE rews) me
   
  rewDeclGroup :: RewMap -> DeclGroup -> M DeclGroup
  rewDeclGroup rews dg =
    case dg of
      NonRecursive d -> NonRecursive <$> rewD rews d
      Recursive ds ->
        do let (leave,rew) = partitionEithers (map consider ds)
               rewGroups   = NE.groupBy sameTParams
                           $ sortBy compareTParams rew
           ds1 <- mapM (rewD rews) leave
           ds2 <- mapM rewSame rewGroups
           return $ Recursive (ds1 ++ concat ds2)
  
    where
    sameTParams    (_,tps1,x,_) (_,tps2,y,_) = tps1 == tps2 && x == y
    compareTParams (_,tps1,x,_) (_,tps2,y,_) = compare (x,tps1) (y,tps2)
  
    consider d   =
      case dDefinition d of
        DPrim         -> Left d
        DForeign _ me ->
          case me of
            Nothing -> Left d
            Just e  -> conExpr e
        DExpr e       -> conExpr e
      where
      conExpr e =
        let (tps,props,e') = splitTParams e
        in if not (null tps) && notFun e'
            then Right (d, tps, props, e')
            else Left d
  
    rewSame ds =
       do new <- forM (NE.toList ds) $ \(d,_,_,e) ->
                   do x <- newName
                      return (d, x, e)
          let (_,tps,props,_) NE.:| _ = ds
              tys            = map (TVar . tpVar) tps
              proofNum       = length props
              addRew (d,x,_) = insertTM (dName d,tys,proofNum) x
              newRews        = foldr addRew rews new
  
          newDs <- forM new $ \(d,newN,e) ->
                     do e1 <- rewE newRews e
                        return ( d
                               , d { dName        = newN
                                   , dSignature   = (dSignature d)
                                           { sVars = [], sProps = [] }
                                   , dDefinition  = DExpr e1
                                   }
                               )
  
          case newDs of
            [(f,f')] ->
                return  [ f { dDefinition =
                                  let newBody = EVar (dName f')
                                      newE = EWhere newBody
                                                [ Recursive [f'] ]
                                  in DExpr $ foldr ETAbs
                                     (foldr EProofAbs newE props) tps
                              }
                        ]
  
            _ -> do tupName <- newTopOrLocalName
                    let (polyDs,monoDs) = unzip newDs
  
                        tupAr  = length monoDs
                        addTPs = flip (foldr ETAbs)     tps
                               . flip (foldr EProofAbs) props
  
                        -- tuple = \{a} p -> (f',g')
                        --                where f' = ...
                        --                      g' = ...
                        tupD = Decl
                          { dName       = tupName
                          , dSignature  =
                              Forall tps props $
                                 TCon (TC (TCTuple tupAr))
                                      (map (sType . dSignature) monoDs)
  
                          , dDefinition =
                              DExpr  $
                              addTPs $
                              EWhere (ETuple [ EVar (dName d) | d <- monoDs ])
                                     [ Recursive monoDs ]
  
                          , dPragmas    = [] -- ?
  
                          , dInfix = False
                          , dFixity = Nothing
                          , dDoc = Nothing
                          }
  
                        mkProof e _ = EProofApp e
  
                        -- f = \{a} (p) -> (tuple @a p). n
  
                        mkFunDef n f =
                          f { dDefinition =
                                DExpr  $
                                addTPs $ ESel ( flip (foldl mkProof) props
                                              $ flip (foldl ETApp) tys
                                              $ EVar tupName
                                              ) (TupleSel n (Just tupAr))
                            }
  
                    return (tupD : zipWith mkFunDef [ 0 .. ] polyDs)
  
  
  --------------------------------------------------------------------------------
  splitTParams :: Expr -> ([TParam], [Prop], Expr)
  splitTParams e = let (tps, e1)   = splitWhile splitTAbs e
                       (props, e2) = splitWhile splitProofAbs e1
                   in (tps,props,e2)
  
  -- returns type instantitaion and how many "proofs" were there
  splitTApp :: Expr -> Int -> (Expr, [Type], Int)
  splitTApp (EProofApp e) n = splitTApp e $! (n + 1)
  splitTApp e0 n            = let (e1,ts) = splitTy e0 []
                              in (e1, ts, n)
    where
    splitTy (ETApp e t) ts = splitTy e (t:ts)
    splitTy e ts           = (e,ts)
  
  notFun :: Expr -> Bool
  notFun (EAbs {})       = False
  notFun (EProofAbs _ e) = notFun e
  notFun _               = True
  
  