content: {-# LANGUAGE PatternGuards, RecordWildCards, ViewPatterns #-}
  {-# OPTIONS_GHC -fno-warn-overlapping-patterns #-}
  
  module Supercompile(supercompiler) where
  
  import Type
  import Simplify
  import Convert
  import Data.Maybe
  
  import qualified Data.Map as Map
  import Control.Monad
  import System.IO.Unsafe
  import System.IO
  import Data.IORef
  import System.Directory
  import Control.Arrow
  import Data.List
  import Debug.Trace
  import Data.Generics.PlateData
  import Control.Monad.State
  import Control.Exception
  
  
  supercompiler :: Prog -> Prog
  supercompiler = supero . simplifyProg . map (unique . normalise) . reduceArity . forwarding
  
  
  supero :: Prog -> Prog
  supero prog = closeLogFile $ runResult $ supercompile func newTerm "supermain" $ toSem $ func "main"
      where func = getFunc prog
  
  
  ---------------------------------------------------------------------
  -- LOGGING
  
  
  {-# NOINLINE logState #-}
  logState :: IORef (Maybe Handle)
  logState = unsafePerformIO $ newIORef Nothing
  
  
  openLogFile :: String -> ()
  openLogFile file = unsafePerformIO $ do
      h <- readIORef logState
      case h of
          Just h -> hClose h
          Nothing -> createDirectoryIfMissing True "log"
      h <- openFile ("log/" ++ file ++ ".log") WriteMode
      writeIORef logState $ Just h
  
  
  appendLogFile :: String -> ()
  appendLogFile msg = unsafePerformIO $ do
      Just h <- readIORef logState
      hPutStrLn h msg
  
  
  closeLogFile :: a -> a
  closeLogFile x = unsafePerformIO $ finally (return $! x) $ do
      h <- readIORef logState
      when (isJust h) $ hClose $ fromJust h
  
  ---------------------------------------------------------------------
  -- TYPE
  
  
  type Env = Var -> Func
  type Tag = String
  type Split = ([Sem], [Fun] -> Expr)
  
  data Sem = Sem [Var] Var (Map.Map Var (Tag,Redex)) deriving (Eq,Ord)
  
  instance Show Sem where
      show (Sem free v mp) = unlines $ line1 : map f (Map.toList mp)
          where line1 = "\\" ++ unwords free ++ " -> " ++ v ++ " where"
                f (k,(s,v)) = "  " ++ k ++ " " ++ show s ++ " = " ++ showExpr (fromRedexShort v)
  
  
  sem :: [Var] -> Var -> Map.Map Var (Tag,Redex) -> Sem
  sem free v mp = relabel $ simplify $ relabel $ Sem free v mp
  
  
  
  -- all the bound redexes must be fully simple
  -- none of the redexes may have a root let statement
  
  data SSimplify = SSimplify {simplifyNew :: Map.Map Var (Tag,Redex)}
  
  simplify :: Sem -> Sem
  simplify (Sem free root mp) = flip evalState (SSimplify Map.empty) $ do
      mapM_ var $ Map.keys mp
      fmap (Sem free root) $ gets simplifyNew
      where
          simp :: Var -> Tag -> Redex -> State SSimplify (Tag, Redex)
          simp v tag x = do
              (tag,x) <- redex tag x
              modify $ \s -> s{simplifyNew = Map.insert v (tag,x) $ simplifyNew s}
              return (tag,x)
  
          var :: Var -> State SSimplify (Tag, Redex)
          var v = do
              new <- gets simplifyNew
              case (Map.lookup v new, Map.lookup v mp) of
                  (Just y, _) -> return y
                  (_, Nothing) -> return $ ("unbound" ++ v,RVar v)
                  (_, Just (tag,x)) -> simp v tag x
  
          redex :: Tag -> Redex -> State SSimplify (Tag, Redex)
          redex tag (RVar v) = var v
          redex tag (RLet vs x) = do
              sequence_ [simp v (tag ++ "_" ++ show i) x | (i,(v,x)) <- zip [1..] vs]
              redex (tag ++ "_0") x
          redex tag o@(RApp v y) = do
              x <- var v
              case x of
                  (tag,RCon c xs) -> return (tag ++ "_", RCon c (xs ++ [y]))
                  _ -> return (tag,o)
          redex tag o@(RCase v alts) = do
              (_,x) <- var v
              case x of
                  RCon{} -> redex (tag++"_") $ head $ concatMap (f x) alts
                  RLit{} -> redex (tag++"_") $ head $ concatMap (f x) alts
                  _ -> return (tag,o)
              where f (RCon c vs) (Patt c2 vs2, x) | c == c2 = [RLet (zip vs2 $ map RVar vs) x]
                    f (RLit c) (PattLit c2, x) | c == c2 = [x]
                    f _ (PattAny, x) = [x]
                    f _ _ = []
          redex tag x = return (tag,x)
  
  
  -- do a GC, variable normalisation, variable flattening
  
  data SRelabel = SRelabel {relabelOld :: Map.Map Var Var, relabelNew :: Map.Map Var (Tag,Redex), relabelVars :: [Var]}
  
  relabel :: Sem -> Sem
  relabel (Sem free root mp) = flip evalState s0 $ do
      root2 <- move root
      old <- gets relabelOld
      free2 <- return [Map.findWithDefault "_" x old | x <- free]
      fmap (Sem free2 root2) $ gets relabelNew
      where
          s0 = SRelabel Map.empty Map.empty ['x':show i | i <- [1..]]
          fresh = do v:vs <- gets relabelVars ; modify $ \s -> s{relabelVars=vs} ; return v
          freshN n = replicateM n fresh
  
          rename x y = modify $ \s -> s{relabelOld = Map.insert x y $ relabelOld s}
          record x y = modify $ \s -> s{relabelNew = Map.insert x y $ relabelNew s}
  
          move v = do
              old <- gets relabelOld
              case (Map.lookup v old, Map.lookup v mp) of
                  (Just y, _) -> return y
                  (_, Nothing) | v `elem` free -> do y <- fresh; rename v y; return y
                               | otherwise -> return v
                  (_, Just (_, RVar y)) -> move y
                  (_, Just (tag,x)) -> do
                      y <- fresh
                      rename v y
                      x <- f Map.empty x
                      record y (tag,x)
                      return y
  
          f mp (RApp x y) = liftM2 RApp (var mp x) (var mp y)
          f mp (RCase x xs) = liftM2 RCase (var mp x) (mapM (alt mp) xs)
          f mp (RFun x) = return $ RFun x
          f mp (RLit x) = return $ RLit x
          f mp (RCon c xs) = liftM (RCon c) $ mapM (var mp) xs
          f mp (RVar x) = liftM RVar (var mp x)
          f mp (RLet vxs x) = do
              let (vs,xs) = unzip vxs
              vs2 <- replicateM (length vs) fresh
              xs2 <- mapM (f mp) xs
              x2 <- f (Map.fromList (zip vs vs2) `Map.union` mp) x
              return $ RLet (zip vs2 xs2) x2
  
          alt mp (Patt c vs, x) = do
              vs2 <- replicateM (length vs) fresh
              x2 <- f (Map.fromList (zip vs vs2) `Map.union` mp) x
              return (Patt c vs2, x2)
          alt mp (v, x) = fmap ((,) v) $ f mp x
  
          var mp v = case Map.lookup v mp of
              Nothing -> move v
              Just y -> return y
  
  
  toSem :: Func -> Sem
  toSem x = sem args "?v" (Map.singleton "?v" ("main",bod))
      where (args,bod) = unique $ toRedex x
  
  
  ---------------------------------------------------------------------
  -- MANAGER
  
  supercompile :: Env -> Term -> Fun -> Sem -> Result ()
  supercompile env t name x | terminate t x = f t name (stop t x)
                            | otherwise = do
          () <- return $ openLogFile name
          () <- return $ appendLogFile $ show t
          let (x',t') = reduce env x
          res <- seen x' ; case res of
              Just y -> addSeen y x >> addResult name (forward y)
              _ -> addSeen name x' >> f ((t += x') += x) name t'
      where
          f t name (xs',gen) = do
              (names,acts) <- mapAndUnzipM (g t) xs'
              () <- addResult name (gen names)
              () <- return $ appendLogFile $ showExpr $ gen names
              sequence_ acts
  
          g t x = do
              res <- seen x ; case res of
                  Just y -> return (y, return ())
                  Nothing -> do
                      [y] <- freshNames 1
                      -- addSeen y x
                      return (y, supercompile env t y x)
  
  
  -- return the Sem, and the split version you'd like to use
  -- both are equivalent
  reduce :: Env -> Sem -> (Sem, Split)
  reduce env = f newTerm
      where
          f t x | () <- appendLogFile $ show x, False = undefined
                | terminate t x = (x, stop t x)
                | Just x' <- step env x = f (t += x) x'
                | otherwise = (x, split x)
  
  
  forward :: Fun -> Expr
  forward = EFun
  
  {-
  
  
  supercompile :: Env -> Term -> Fun -> Sem -> Result ()
  supercompile env global name x = logFile ("func_" ++ name) $ \out -> f out newTerm x
      where
          f out local x | () <- out ("before step\n" ++ show x), False = error "impossible"
                        | Just e <- terminate local x = g out (Just e) [] x
                        | Just x2 <- step env x = f out (local+=x) x2
                        | otherwise = g out Nothing [] x
  
          g out evidence vars x = do
              () <- return $ out ("before split\n" ++ show x)
              r <- seen x
              case r of
                  Just y -> addResult name (EFun y)
                  Nothing -> do
                      addSeen name x
                      let (y,ys) = (id *** (++ vars)) $ split evidence x
                      case terminate global y of
                          Just e -> do
                              g out (Just e) ys y
                          Nothing -> do
                              () <- return $ out ("after split\n" ++ show y)
                              addSeen name y
                              vs <- freshNames $ length ys
                              addResult name $ apply y vs
                              () <- return $ out ("answer\n" ++ showExpr (apply y vs))
                              zipWithM_ (supercompile env (global+=y)) vs ys
  
  -}
  
  ---------------------------------------------------------------------
  -- RESULTS
  
  type Result a = State ResultState a
  
  data ResultState = ResultState {resultSeen :: [(Sem,Fun)], resultFresh :: !Int, resultProg :: Prog}
  
  
  runResult :: Result a -> Prog
  runResult x = unsafePerformIO $ do
      writeFile "log/result.log" ""
      return $ resultProg $ execState x $ ResultState [] 1 []
  
  addSeen :: Fun -> Sem -> Result ()
  addSeen fun sem = modify $ \r -> r{resultSeen = (sem,fun) : resultSeen r}
  
  addResult :: Fun -> Expr -> Result ()
  addResult fun expr = unsafePerformIO $ do
      let x = Func fun [] expr
      appendFile "log/result.log" (showFunc x ++ "\n\n")
      return $ modify $ \r -> r{resultProg = x : resultProg r}
  
  freshNames :: Int -> Result [Fun]
  freshNames n = do
      r <- get
      put r{resultFresh = resultFresh r + n}
      return $ map ((:) 'f' . show) $ take n [resultFresh r ..]
  
  seen :: Sem -> Result (Maybe Fun)
  seen sem = fmap (lookup sem) $ gets resultSeen
  
  
  ---------------------------------------------------------------------
  -- TERMINATION
  
  data Term = Term [Past] deriving Show
  
  type Past = [(String,Int)]
  
  past (Sem _ _ mp) = map (head &&& length) $ sort $ group $ sort $ map fst $ Map.elems mp
  
  -- return True if new is greater than old in every way
  bad ((x1,x2):xs) ((y1,y2):ys) = case compare x1 y1 of
      LT -> bad xs ((y1,y2):ys)
      GT -> False
      EQ -> x2 >= y2 && bad xs ys
  bad _ [] = True
  bad [] _ = False
  
  
  newTerm :: Term
  newTerm = Term []
  
  terminate :: Term -> Sem -> Bool
  terminate (Term xs) sem | any (bad $ past sem) xs = True -- error $ show (sem, xs)
                          | otherwise = False
  
  (+=) :: Term -> Sem -> Term
  (+=) (Term xs) sem = Term $ past sem : xs
  
  
  
  stop :: Term -> Sem -> Split
  stop term (Sem free root mp) = trace "stop" $ (map snd ss, \names -> eLams free $ eLet (zip keep $ zipWith g names ss) (eVar root))
      where
          keep = fixEq op $ Map.keys mp
          ss = map (f keep) keep
          f keep w = semTop (delete w $ free ++ keep) w (Map.filterWithKey (\k _ -> k `notElem` keep || k == w) mp)
          g name (vs,_) = eApps (eFun name) (map eVar vs)
  
          op xs = h xs poss
              where
                  poss = filter (\x -> (<=1) $ length $ filter (elem x . nub . childrenBi . flip Map.lookup mp) $ delete x xs) $ delete root xs
                  h xs [] = xs
                  h xs (p:oss) | all (not . terminate term . snd . f k2) k2 = h k2 oss
                               | otherwise = h xs oss
                      where k2 = delete p xs
  
  
  ---------------------------------------------------------------------
  -- SPLIT
  
  
  split :: Sem -> Split
  split (Sem free root mp) = (a, eLams free . b)
      where
          (a,b) = f root
      
          f v = case fmap snd $ Map.lookup v mp of
              Just (RCase w alt) | w `Map.member` mp -> f w
                                 | otherwise -> splitCase v w alt
              Just (RCon c xs) -> splitApp v (eCon c) xs
              Just (RApp x y) | x `Map.member` mp -> f x
                              | otherwise -> splitApp v (eFun x) [y]
              Just (RFun v) -> splitFun
              Just (RLit x) -> ([], \_ -> eLit x)
              Nothing -> ([], \_ -> eVar v)
              x -> ([], const $ eFun $ "how do i split on " ++ show x)
  
  
          splitFun = ([b], \[name] -> eLam "newArg" $ eApps (eFun name) $ map eVar a)
              where (a,b) = semTop (free ++ ["newArg"]) "newRoot" $ Map.insert "newRoot" ("newRoot",RApp root "newArg") mp
  
  
          splitApp v gen xs = (map snd ss, \names -> eLet ((v, eApps gen (map eVar xs)) :  zip (delete v keep) (zipWith g names ss)) (eVar root))
              where
                  keep = share (Sem free root mp) $ root:v:xs
                  ss = map f $ delete v keep
                  f w = semTop (delete w $ free ++ keep) w (Map.filterWithKey (\k _ -> k `notElem` keep || k == w) mp)
                  g name (vs,_) = eApps (eFun name) (map eVar vs)
  
  
          -- type Split = ([Sem], [Fun] -> Expr)
          splitCase v on alt = (map snd ss,
                  \names -> eCase (eVar on) $ zipWith3 g names ss alt)
              where
                  ss = map f alt
                  f (Patt c vs, x) = semTop (vs ++ delete on free) root $ Map.insert on ("con",RCon c vs) mp
                  f (PattLit c, x) = semTop (delete on free) root $ Map.insert on ("lit",RLit c) mp
                  f (PattAny, x) = semTop free root $ Map.insert v ("case",x) mp
                  g name (vs,_) (pat,_) = (pat, eApps (eFun name) (map eVar vs))
  
  
  -- top level sem, which can throw away free vars
  semTop :: [Var] -> Var -> Map.Map Var (Tag,Redex) -> ([Var], Sem)
  semTop free root mp = ([v | (v,w) <- zip free free2, w /= "_"], Sem (filter (/= "_") free2) root2 mp2)
      where Sem free2 root2 mp2 = sem free root mp
  
  
  -- given a Sem, I need these variables at the top level, which other ones to keep sharing? only those from mp
  -- at worst, will be all the bound variables (never puts free vars in)
  -- at best, will be just those in the input set that are bound in mp
  share :: Sem -> [Var] -> [Var]
  share (Sem free root mp) top = fixEq f (top \\ free)
      where
          f xs = filter (`Map.member` mp) $ sort $ nub $ xs ++ new
              where new = map head $ filter ((>1) . length) $ group $ sort $ concatMap follow xs
          
          follow x = case Map.lookup x mp of
              Nothing -> []
              Just r -> childrenBi r :: [String]
  
  
  fixEq f x = if x == x2 then x2 else fixEq f x2
      where x2 = f x
  
  
  ---------------------------------------------------------------------
  -- STEP
  
  
  step :: Env -> Sem -> Maybe Sem
  step e (Sem free root mp) = f [] root
      where
          f app v = case fmap snd $ Map.lookup v mp of
              Nothing -> Nothing
              Just (RCase on _) -> f [] on
              Just (RApp x y) -> f ((v,y):app) x
              Just (RFun x) | length args == 0 -> Just $ sem free root $ Map.insert v (x,bod) mp
                            | length args <= length app -> Just $ sem free root $ Map.insert (fst $ last_ "step" app2) (x,new) mp
                  where new = RLet (zip args $ map (RVar . snd) app2) bod
                        app2 = take (length args) app
                        (args,bod) = unique $ toRedex $ e x
              _ -> Nothing
  
  
  step e sem@(Sem free root mp) = f $ flatten sem
      where
          f ((v, Just (_, RFun name)):rest) = g v [] rest
              where
                  (args,bod) = unique $ toRedex $ e name
  
                  g v seen _ | length seen == length args = Just $ sem free root $ Map.insert v (name,bod2) mp
                      where bod2 = RLet (zip args $ map RVar seen) bod
                  g _ seen ((v, Just (_, RApp _ w)):rest) = g fun v (w:seen) 
  
          g (args,bod) v seen _ | length seen == length args = Just $ sem free root $ Map.insert v (v,bod2) mp
          
          RLet replace sem val
          g fun _ seen ((v, Just (_, RApp _ w)):rest) = g fun v (w:seen) rest
          g _ _ _ _ = Nothing
  
              where
                  (args,bod) = unique $ toRedex $ e x
                  (rep,app) = g v [] rest
  
                  g v seen _ | length seen == length args = (v, seen)
                  g v seen (v, Just 
  
      
      
          f app v = case fmap snd $ Map.lookup v mp of
              Nothing -> Nothing
              Just (RCase on _) -> f [] on
              Just (RApp x y) -> f ((v,y):app) x
              Just (RFun x) | length args == 0 -> Just $ sem free root $ Map.insert v (x,bod) mp
                            | length args <= length app -> Just $ sem free root $ Map.insert (fst $ last_ "step" app2) (x,new) mp
                  where new = RLet (zip args $ map (RVar . snd) app2) bod
                        app2 = take (length args) app
                        (args,bod) = unique $ toRedex $ e x
              _ -> Nothing
  
  
  
  
  last_ msg [] = error $ "last with " ++ msg
  last_ msg x = last x
  
  
  
  -- the first element may be a nothing, all the others must be a Just
  flatten :: Sem -> [(Var, Maybe (String,Redex))]
  flatten (Sem free root mp) = reverse $ f root
      where
          f v = case Map.lookup v mp of
              Nothing -> [(v, Nothing)]
              Just (s,x) -> (v, Just (s,x)) : case force x of
                  Nothing -> []
                  Just v -> f v
  
  
  -- if you force this redex, which var will you force next
  force :: Redex -> Maybe Var
  force (RVar v) = Just v
  force (RApp v _) = Just v
  force (RCase v _) = Just v
  force _ = Nothing
  