content: {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE Rank2Types #-}
  {-# LANGUAGE PatternGuards #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  
  --------------------------------------------------------------------
  -- |
  -- Copyright :  (c) Edward Kmett and Dan Doel 2014
  -- License   :  BSD2
  -- Maintainer:  Dan Doel <dan.doel@gmail.com>
  -- Stability :  experimental
  -- Portability: non-portable
  --------------------------------------------------------------------
  
  module Ermine.Core.Optimizer
    ( optimize
    , rewriteCore
    , rewriteCoreDown
    ) where
  
  import Bound
  import Bound.Var
  import Bound.Scope
  import Control.Applicative
  import Control.Lens
  import Control.Monad.Writer
  import Data.List (genericLength, genericSplitAt, genericIndex)
  import Data.Traversable (sequenceA)
  import Data.Word
  import Ermine.Syntax
  import Ermine.Syntax.Convention
  import Ermine.Syntax.Core
  import Ermine.Syntax.Scope
  import Ermine.Unification.Sharing
  
  -- | Optimize core expressions by alternating between the different optimization passes several times.
  optimize :: Core Convention c -> Core Convention c
  optimize c = runIdentity . runSharing c $ optimize' 10 c
  
  optimize' :: (Applicative m, MonadWriter Any m)
            => Int -> Core Convention c -> m (Core Convention c)
  optimize' 0 c = return c
  optimize' n c = do (c', Any b) <- listen $ suite c
                     if b then optimize' (n-1) c' else return c
   where
   suite = rewriteCoreDown lamlam
       >=> rewriteCore betaVar
       >=> rewriteCoreDown specCase
       >=> rewriteCore etaDict
  
  rewriteCoreDown :: forall m c cc. (Applicative m, MonadWriter Any m)
                  => (forall d. Core cc d -> m (Core cc d)) -> Core cc c -> m (Core cc c)
  rewriteCoreDown opt = go
   where
   go :: forall e. Core cc e -> m (Core cc e)
   go c = sharing c (opt c) >>= \ xs -> case xs of
     l@(Lam cc e)         -> sharing l $ Lam cc <$> goS e
     d@(Data cc n g l)    -> sharing d $ Data cc n g <$> traverse go l
     d@(Prim cc r g l)    -> sharing d $ Prim cc r g <$> traverse go l
     a@(App cc f x)       -> sharing a $ App cc <$> go f <*> go x
     l@(Let d b)          -> sharing l $ Let <$> sharing d (traverse goS d) <*> goS b
     s@(Case e b d)       -> sharing s $ Case <$> go e <*> sharing b ((traverse.matchBody) goS b) <*> sharing d (traverse goS d)
     d@(Dict su sl)       -> sharing d $ Dict <$> sharing su (traverse go su) <*> sharing sl (traverse goS sl)
     x@HardCore{}         -> return x
     x@Var{}              -> return x
   goS :: forall b e. Scope b (Core cc) e -> m (Scope b (Core cc) e)
   goS s = sharing s . inScope go $ s
  
  rewriteCore :: forall m c cc. (Applicative m, MonadWriter Any m)
              => (forall d. Core cc d -> m (Core cc d)) -> Core cc c -> m (Core cc c)
  rewriteCore opt = go
   where
   go :: forall e. Core cc e -> m (Core cc e)
   go c = sharing c $ opt =<< case c of
     l@(Lam cc e)         -> sharing l $ Lam cc <$> goS e
     d@(Data cc n g l)    -> sharing d $ Data cc n g <$> traverse go l
     d@(Prim cc r g l)    -> sharing d $ Prim cc r g <$> traverse go l
     a@(App cc f x)       -> sharing a $ App cc <$> go f <*> go x
     l@(Let d b)          -> sharing l $ Let <$> sharing d (traverse goS d) <*> goS b
     s@(Case e b d)       -> sharing s $ Case <$> go e <*> sharing b ((traverse.matchBody) goS b) <*> sharing d (traverse goS d)
     d@(Dict su sl)       -> sharing d $ Dict <$> sharing su (traverse go su) <*> sharing sl (traverse goS sl)
     x@HardCore{}         -> return x
     x@Var{}              -> return x
   goS :: forall b e. Scope b (Core cc) e -> m (Scope b (Core cc) e)
   goS s = sharing s . inScope go $ s
  
  -- | Turns @\{x..} -> \{y..} -> ...@ into @\{x.. y..} -> ...@ for all lambda variants
  lamlam :: forall c cc m. (Functor m, MonadWriter Any m) => Core cc c -> m (Core cc c)
  lamlam (Lam cc0 e0) = slurp False cc0 (fromScope e0)
   where
   slurp :: forall e. Bool -> [cc] -> Core cc (Var Word64 e) -> m (Core cc e)
   slurp _ m (Lam n b) | j <- fromIntegral $ length m = slurp True (m ++ n) (instantiate (\i -> pure $ B $ j + i) b)
   slurp b m c = Lam m (toScope c) <$ tell (Any b)
  lamlam c = return c
  
  -- | 'Lam D' is strict, so η-reduction for it is sound. η-reduce.
  --
  -- Todo: generalize this to larger lambdas and also extend it to cover other strict lambdas like U and N
  etaDict :: forall c m. (Functor m, MonadWriter Any m) => Core Convention c -> m (Core Convention c)
  etaDict c@(Lam [D] (Scope (App D f (Var (B 0))))) = case sequenceA f of
    B _ -> return c
    F g -> join g <$ tell (Any True)
  etaDict c = return c
  
  -- | β-reduces redexes like @(\x.. -> e) v..@ where @v..@ is all variables for all calling conventions
  betaVar :: forall c m. (Applicative m, MonadWriter Any m) => Core Convention c -> m (Core Convention c)
  betaVar = collapse []
   where
   collapse stk (App C f x)
     | has _Var x || has _Slot x || has _Super x || has _Lit x = collapse (x:stk) f
   collapse stk (Lam cc body@(Scope body'))
     | len < n = do
       tell (Any True)
       let replace i
             | i < len   = F $ stk `genericIndex` i
             | otherwise = B $ i - len
       return $ Lam (drop (fromIntegral len) cc) $ Scope $ fmap (unvar replace F) body'
     | (args, stk') <- genericSplitAt n stk = do
       tell (Any True)
       collapse stk' $ instantiate (genericIndex args) body
     where
       len = genericLength stk
       n   = genericLength cc
   collapse stk c = return $ apps c stk
  
  -- | Specializes a case expression to a known constructor.
  --
  -- TODO: switch to
  --
  -- (\\ a b c -> let x = Foo a b c in body) arg1 .. argN
  --
  -- and rely on beta reduction and other checks for work-safe application reduction
  --
  -- This has the benefit that it works for all calling conventions
  specCase :: forall m c. (Applicative m, MonadWriter Any m) => Core Convention c -> m (Core Convention c)
  specCase c@(Case dat@(Data cc tg g as) bs d)
    -- TODO: Use a LamHash around this for all the unboxed arguments, and AppHash each to an argument.
    | any (/=C) cc = pure c -- until we change this to use Lam as described above
    | Just (Match _cc _g body) <- bs ^. at tg =
      Let ((Scope . Data cc tg g $ pure . B <$> [1..fromIntegral arity]) : map lift as)
          (mapBound fromIntegral body)
        <$ tell (Any True)
    | Just e <- d = Let [lift dat] (mapBound (const 0) e) <$ tell (Any True)
    | otherwise = error "PANIC: non-exhaustive case with no default"
    where arity = length cc
  specCase c = pure c
  