@@ -1,6 +1,9 @@
1+{-# LANGUAGE PatternGuards, DeriveFunctor #-}
2+
13module IRTS.Lang where
24
35import Control.Monad.State hiding (lift)
6+import Control.Applicative hiding (Const)
47import Idris.Core.TT
58import Idris.Core.CaseTree
69import Data.List
@@ -92,10 +95,13 @@ data FType = FArith ArithTy
9295 | FAny
9396 deriving (Show, Eq)
9497
95−data LAlt = LConCase Int Name [Name] LExp
96− | LConstCase Const LExp
97− | LDefaultCase LExp
98− deriving (Show, Eq)
98+-- FIXME: Why not use this for all the IRs now?
99+data LAlt' e = LConCase Int Name [Name] e
100+ | LConstCase Const e
101+ | LDefaultCase e
102+ deriving (Show, Eq, Functor)
103+
104+type LAlt = LAlt' LExp
99105
100106data LDecl = LFun [LOpt] Name [Name] LExp -- options, name, arg names, def
101107 | LConstructor Name Int Int -- constructor name, tag, arity
@@ -125,9 +131,9 @@ liftAll :: [(Name, LDecl)] -> [(Name, LDecl)]
125131liftAll xs = concatMap (\ (x, d) -> lambdaLift x d) xs
126132
127133lambdaLift :: Name -> LDecl -> [(Name, LDecl)]
128−lambdaLift n (LFun _ _ args e)
134+lambdaLift n (LFun opts _ args e)
129135 = let (e', (LS _ _ decls)) = runState (lift args e) (LS n 0 []) in
130− (n, LFun [] n args e') : decls
136+ (n, LFun opts n args e') : decls
131137lambdaLift n x = [(n, x)]
132138
133139getNextName :: State LiftState Name
@@ -196,6 +202,98 @@ lift env LNothing = return $ LNothing
196202allocUnique :: LDecl -> LDecl
197203allocUnique x = x
198204
205+inlineAll :: [(Name, LDecl)] -> [(Name, LDecl)]
206+inlineAll lds = let defs = addAlist lds emptyContext in
207+ map (\ (n, def) -> (n, doInline defs def)) lds
208+
209+nextN :: State Int Name
210+nextN = do i <- get
211+ put (i + 1)
212+ return $ sMN i "in"
213+
214+-- Inline inside a declaration. Variables are still Name at this stage.
215+-- Need to preserve uniqueness of variable names in the resulting definition,
216+-- so invent a new name for every variable we encounter
217+doInline :: LDefs -> LDecl -> LDecl
218+doInline defs d@(LConstructor _ _ _) = d
219+doInline defs (LFun opts topn args exp)
220+ = let res = evalState (inlineWith (map (\n -> (n, LV (Glob n))) args) exp) 0 in
221+ LFun opts topn args res
222+ where
223+ inlineWith :: [(Name, LExp)] -> LExp -> State Int LExp
224+ inlineWith env var@(LV (Glob n)) = case lookup n env of
225+ Just t -> return t
226+ Nothing -> return var
227+ inlineWith env (LLazyApp n es) = LLazyApp n <$> (mapM (inlineWith env) es)
228+ inlineWith env (LForce e) = LForce <$> inlineWith env e
229+ inlineWith env (LLazyExp e) = LLazyExp <$> inlineWith env e
230+ -- Extend the environment for Let and Lam so that bound names aren't
231+ -- expanded with any top level argument definitions they shadow
232+ inlineWith env (LLet n val sc)
233+ = do n' <- nextN
234+ LLet n' <$> inlineWith env val <*>
235+ inlineWith ((n, LV (Glob n')) : env) sc
236+ inlineWith env (LLam args sc)
237+ = do ns' <- mapM (\n -> do n' <- nextN
238+ return (n, n')) args
239+ LLam (map snd ns') <$>
240+ inlineWith (map (\ (n,n') -> (n, LV (Glob n'))) ns' ++ env) sc
241+ inlineWith env (LProj exp i) = LProj <$> inlineWith env exp <*> return i
242+ inlineWith env (LCon loc i n es)
243+ = LCon loc i n <$> mapM (inlineWith env) es
244+ inlineWith env (LCase ty e alts)
245+ = LCase ty <$> inlineWith env e <*> mapM (inlineWithAlt env) alts
246+ inlineWith env (LOp f es) = LOp f <$> mapM (inlineWith env) es
247+ -- the interesting case!
248+ inlineWith env (LApp t (LV (Glob n)) es)
249+ | [LFun opts _ args body] <- lookupCtxt n defs,
250+ Inline `elem` opts,
251+ length es == length args
252+ = do es' <- mapM (inlineWith env) es
253+ inlineWith (zip args es' ++ env) body
254+-- inlineWith
255+-- -- inlineWith env
256+-- (substLExp (zip args es) body)
257+ inlineWith env (LApp t f es)
258+ = LApp t <$> inlineWith env f <*> mapM (inlineWith env) es
259+ inlineWith env (LForeign l t s args)
260+ = LForeign l t s <$> mapM (\(t, e) -> do e' <- inlineWith env e
261+ return (t, e')) args
262+ inlineWith env t = return t
263+
264+ inlineWithAlt env (LConCase i n es rhs)
265+ = do ns' <- mapM (\n -> do n' <- nextN
266+ return (n, n')) es
267+ LConCase i n (map snd ns') <$>
268+ inlineWith (map (\ (n,n') -> (n, LV (Glob n'))) ns' ++ env) rhs
269+ inlineWithAlt env (LConstCase c e) = LConstCase c <$> inlineWith env e
270+ inlineWithAlt env (LDefaultCase e) = LDefaultCase <$> inlineWith env e
271+
272+-- substLExp :: [(Name, LExp)] -> LExp -> LExp
273+-- substLExp env (LV (Glob n)) | Just t <- lookup n env = t
274+-- substLExp env (LApp t f as) = LApp t (substLExp env f) (map (substLExp env) as)
275+-- substLExp env (LLazyApp n as) = LLazyApp n (map (substLExp env) as)
276+-- substLExp env (LLazyExp e) = LLazyExp (substLExp env e)
277+-- substLExp env (LForce e) = LForce (substLExp env e)
278+-- substLExp env (LLet n v sc) = LLet n (substLExp env v)
279+-- (substLExp ((n, LV (Glob n)) : env) sc)
280+-- substLExp env (LLam args sc)
281+-- = LLam args (substLExp (map (\n -> (n, LV (Glob n))) args ++ env) sc)
282+-- substLExp env (LProj e i) = LProj (substLExp env e) i
283+-- substLExp env (LCon up i n es) = LCon up i n (map (substLExp env) es)
284+-- substLExp env (LCase ty e alts) = LCase ty (substLExp env e) (map substC alts)
285+-- where
286+-- substC (LConCase i n args e)
287+-- = LConCase i n args
288+-- (substLExp (map (\n -> (n, LV (Glob n))) args ++ env) e)
289+-- substC (LConstCase i e) = LConstCase i (substLExp env e)
290+-- substC (LDefaultCase e) = LDefaultCase (substLExp env e)
291+-- substLExp env (LForeign l t s args)
292+-- = LForeign l t s (map (\ (ty, e) -> (ty, substLExp env e)) args)
293+-- substLExp env (LOp fn es) = LOp fn (map (substLExp env) es)
294+-- substLExp env t = t
295+
296+
199297-- Return variables in list which are used in the expression
200298
201299usedArg env n | n `elem` env = [n]