content: {-# LANGUAGE LambdaCase, RecordWildCards #-}
  module Lambda.Util where
  
  import Data.Functor.Foldable
  import qualified Data.Set as Set
  import qualified Data.Foldable
  import qualified Data.Map as Map
  import Data.List (unzip6)
  
  import Transformations.Names
  import Transformations.Util
  import Lambda.Syntax
  
  mkNameEnv :: Exp -> NameEnv
  mkNameEnv exp = NameEnv mempty (cata folder exp) where
    folder e = foldNameDefExpF Set.singleton e `mappend` Data.Foldable.fold e
  
  foldLocalNameDefExp :: (Monoid m) => (Name -> m) -> Exp -> m
  foldLocalNameDefExp f = foldLocalNameDefExpF f . project
  
  foldLocalNameDefExpF :: (Monoid m) => (Name -> m) -> ExpF a -> m
  foldLocalNameDefExpF f = \case
    DefF _ args _ -> mconcat $ map (f . fst) args
    ProgramF{}    -> mempty
    e             -> foldNameDefExpF f e
  
  foldNameDefExpF :: (Monoid m) => (Name -> m) -> ExpF a -> m
  foldNameDefExpF f = \case
    ProgramF{..}            -> mconcat $ map (f . sName) pStaticDataF
    DefF n args _           -> mconcat $ f n : map (f . fst) args
    LetF bs _               -> mconcat $ map (f . fst3) bs
    LetRecF bs _            -> mconcat $ map (f . fst3) bs
    LetSF bs _              -> mconcat $ map (f . fst3) bs
    AltF a (NodePat _ bs) _ -> mconcat $ f a : map f bs
    ClosureF _ bs _         -> mconcat $ map (f . fst) bs
    _                       -> mempty
  
  fst3 :: (a, b, c) -> a
  fst3 (a, _, _) = a
  
  mapNameExp :: (Name -> Name) -> Exp -> Exp
  mapNameExp f = \case
    prg@Program{..} -> prg { pPublicNames = map f pPublicNames
                           , pForeignExportedNames = map f pForeignExportedNames
                           , pStaticData  = [sd {sName = f $ sName sd} | sd <- pStaticData]
                           }
    Def n args e    -> Def (f n) [(f a, t) | (a, t) <- args] e
    exp             -> mapLocalNameExp f exp
  
  mapLocalNameExp :: (Name -> Name) -> Exp -> Exp
  mapLocalNameExp f = \case
    -- def
    Def n args e            -> Def n [(f n, t) | (n, t) <- args] e -- NOTE: do not touch top level function name
    LetRec bs e             -> LetRec [(f n, t, a) | (n, t, a) <- bs] e
    Let bs e                -> Let  [(f n, t, a) | (n, t, a) <- bs] e
    LetS bs e               -> LetS [(f n, t, a) | (n, t, a) <- bs] e
    Alt a (NodePat n bs) e  -> Alt (f a) (NodePat n (map f bs)) e
    Closure vs bs e         -> Closure (map f vs) [(f n, t) | (n, t) <- bs] e
    -- use
    App n l                 -> App (f n) (map f l)
    Case n a                -> Case (f n) a
    Con n l                 -> Con (f n) (map f l)
    Var n                   -> Var (f n)
    -- other
    exp                     -> exp
  
  isAtom :: Exp -> Bool
  isAtom = \case
    Var{} -> True
    Lit{} -> True
    _     -> False
  
  -- merges let seqences of the same kind
  smashLet :: Exp -> Exp
  smashLet = cata folder where
    folder = \case
      LetF l1 (Let l2 e)    -> Let (l1 ++ l2) e
      LetSF l1 (LetS l2 e)  -> LetS (l1 ++ l2) e
      e                     -> embed e
  
  mapNameTy :: (Name -> Name) -> Ty -> Ty
  mapNameTy f = \case
    TyCon n c l   -> TyCon (f n) c l
    TyVar n       -> TyVar (f n)
    TySimple n t  -> TySimple (f n) t
    TyFun n t l   -> TyFun (f n) t l
  
  foldNameTyF :: (Monoid m) => (Name -> m) -> TyF a -> m
  foldNameTyF f = \case
    TyConF n _ _  -> f n
    TyVarF n      -> f n
    TySimpleF n _ -> f n
    TyFunF n _ _  -> f n
  
  -- TODO: validate merge operation (i.e. con group and external definiton must match if names are matching, def names must be unique)
  concatPrograms :: [Program] -> Program
  concatPrograms prgs = prg where
    (exts, cgroups, publics, fexported, sdata, defs) = unzip6 [(e, c, p, f, s, d) | Program e c p f s d <- prgs]
    nubExts l       = Map.elems $ Map.fromList [(eName x, x) | x <- l]
    nubConGroups l  = Map.elems $ Map.fromList [(cgName x, x) | x <- l]
    prg = Program
      { pExternals    = nubExts $ concat exts
      , pConstructors = nubConGroups $ concat cgroups
      , pPublicNames  = Set.toList . Set.fromList $ concat publics
      , pForeignExportedNames = Set.toList . Set.fromList $ concat fexported
      , pStaticData   = concat sdata
      , pDefinitions  = concat defs
      }
  
  sortProgramDefs :: Program -> Program
  sortProgramDefs prg@Program{..} = prg
    { pConstructors = Map.elems $ Map.fromList [(cgName c, c) | c <- pConstructors]
    , pPublicNames  = Set.toList $ Set.fromList pPublicNames
    , pForeignExportedNames = Set.toList $ Set.fromList pForeignExportedNames
    , pStaticData   = Map.elems $ Map.fromList [(sName d, d) | d <- pStaticData]
    , pDefinitions  = Map.elems $ Map.fromList [(n,d) | d@(Def n _ _) <- pDefinitions]
    }
  