content: {-# LANGUAGE PatternGuards #-}
  
  module Simplify(reduceArity, simplifyProg, simplifyExpr, forwarding) where
  
  import Type
  import Data.Generics.PlateData
  import Data.List
  import qualified Data.Map as Map
  
  
  -- perform some simplification cleanups
  -- run both before and after supercompilation
  
  -- global:
  -- foo x y z = bar x y z, replace all foo with bar, transitively
  -- foo x = .... ..., where x is never used. Eliminate foo the argument where possible
  
  -- let x = y in z, where x is used once = z[x/y]
  
  
  simplifyProg :: Prog -> Prog
  simplifyProg = descendBi (unique . simplifyExpr)
  
  
  simplifyExpr :: Expr -> Unique Expr
  simplifyExpr = transformM f
      where
          f (ELet _ xy z) | not $ null sub = f . eLet keep =<< rep sub z
              where (sub,keep) = partition (\(x,y) -> linear x z || simple y) xy
  
          -- only safe because variables are unique
          f (ECase _ (ELet _ xy z) alts) = f . eLet xy =<< f (eCase z alts)
  
          f (ECase _ on alts) | (PattAny,ECase _ on2 alts2) <- last alts, on == on2 = f $ eCase on (init alts ++ alts2)
  
          f (ECase _ (EVar on) alts) | any fst alts2 = fmap (eCase (EVar on)) $ sequence $ map snd alts2
              where alts2 = map g alts
                    g (Patt c vs, x) = let b = any (`Map.member` getI x) vs in (b, fmap ((,) (Patt c vs)) $ rep [(on,eApps (eCon c) (map eVar vs)) | b] x)
                    g x = (False, return x)
  
          f (ECase _ on alts) | (ECon c, vs) <- fromEApps on = head $ concatMap (g c vs) alts ++ [error "simplifyExpr bougus case"]
              where g c vs (PattAny,x) = [return x]
                    g c vs (Patt c2 vs2, x) = [rep (zip vs2 vs) x | c == c2]
                    g c vs _ = []
  
          f (ECase _ on [(Patt "(:)" vs,x), (PattAny,y)]) = f $ eCase on [(Patt "(:)" vs,x), (Patt "[]" [],y)]
  
          f x = return x
  
  
          rep [] z = return z
          rep xy z = substsWith f (Map.fromList xy) z
  
  
  simple EVar{} = True
  simple EFun{} = True
  simple ECon{} = True
  simple _ = False
  
  
  reduceArity :: Prog -> Prog
  reduceArity prog = deleteArgs (Map.keys $ reduce mp) prog
      where
          mn = minArity prog
          mp = Map.fromList [((a,i), nub v) | (a,b) <- Map.toList mn, let fun = getFunc prog a, i <- [0..b-1],
                             i < length (funcArgs fun), Just v <- [callers (funcArgs fun !! i) (funcBody fun)]]
  
          reduce x = if Map.size x == Map.size x2 then x else reduce x2
              where x2 = Map.filter (\xs -> all (`Map.member` x) xs) x
  
  
  deleteArgs :: [(Fun,Int)] -> Prog -> Prog
  deleteArgs xs = map f
      where
          mp = Map.unionsWith (++) [Map.singleton a [b] | (a,b) <- xs]
  
          f (Func name args bod) = Func name (del name args) $ g bod
  
          g x | (EFun y, z) <- fromEApps x = eApps (EFun y) (del y $ map (descend g) z)
          g x = descend g x
  
          del name xs = [x | (i,x) <- zip [0..] xs, i `notElem` ys]
              where ys = Map.findWithDefault [] name mp
  
  
  -- what is the minimum calling arity of each function
  minArity :: Prog -> Map.Map Fun Int
  minArity prog = Map.filterWithKey (\k v -> v /= 0 && k `elem` map funcName prog) $ Map.unionsWith min $ map f $ childrenBi prog
      where
          f x = Map.unionsWith min $ case fromEApps x of
              (EFun a, b) -> Map.singleton a (length b) : map f (concatMap children b)
              _ -> map f $ children x
  
  
  -- which functions get given this argument, Nothing if it's used directly
  callers :: Var -> Expr -> Maybe [(Fun,Int)]
  callers v x | (EFun y, z) <- fromEApps x = comb $
      Just (map ((,) y) $ findIndices (== EVar v) z) :
      map (callers v) (concatMap children z)
  
  callers v (EVar y) = if v == y then Nothing else Just []
  
  callers v (ELet _ xy z) = comb $
      (if v `elem` map fst xy then Just [] else callers v z) :
      (map (callers v . snd) xy)
  
  callers v (ECase _ on alts) = comb $
      (callers v on) :
      [if v `elem` pattVars a then Just [] else callers v b | (a,b) <- alts]
  
  callers v x = comb $ map (callers v) $ children x
  
  comb [] = Just []
  comb (x:xs) = both x $ comb xs
  
  both (Just a) (Just b) = Just $ a ++ b
  both _ _ = Nothing
  
  
  
  -- if a is just a synonym for b, replace every occurence of a with b
  -- doesn't work if multiple functions are just _|_
  forwarding  :: Prog -> Prog
  forwarding prog = filter (not . flip Map.member fwd . funcName) $ transformBi f prog
      where
          f (EFun x) = case Map.lookup x fwd of
                           Nothing -> EFun x
                           Just y -> f $ EFun y
          f x = x
  
          fwd = Map.fromList $ concatMap g prog
          g (Func name args bod) | (EFun x,y) <- fromEApps bod, map eVar args == y, name /= "main" = [(name,x)]
          g _ = []
  