24f580d · Idris-dev · git.iko.soy

git.iko.soy/Idris-dev

Remove a few more lambdas in Lang IR

Slight tweak to lambda lifting, combining nested lambdas. Also eta contraction to change \x -> f x to f, which helps when 'f' is a local variable so no need to construct a new thunk.

commit24f580d45455b3fee35d0e96e48415612e58aaed authorEdwin Brady <ecb10@st-andrews.ac.uk> dateFri 15 Dec 2017, 14:09 (9 years ago) parent7282c0a

2 files changed+55−33

compare with master

unifiedsplit

src/IRTS/Lang.hs+9−8 src/IRTS/LangOpts.hs+46−25

src/IRTS/Lang.hs+9−8

@@ -15,7 +15,7 @@ import Idris.Core.TT

1515

1616import Control.Monad.State hiding (lift)

1717import Data.Data (Data)

18−import Data.List

18+import Data.List

1919import qualified Data.Map.Strict as Map

2020import Data.Typeable (Typeable)

2121import GHC.Generics (Generic)

@@ -141,7 +141,7 @@ addTags i ds = tag i ds []

141141 tag i (x : as) acc = tag i as (x : acc)

142142 tag i [] acc = (i, reverse acc)

143143

144−data LiftState = LS (Maybe Name) Int [(Name, LDecl)]

144+data LiftState = LS (Maybe Name) Int [(Name, LDecl)]

145145 (Map.Map ([Name], LExp) Name) -- map from args/expressions

146146 -- to names, so we don't create the same function

147147 -- multiple times

@@ -171,12 +171,12 @@ renameArgs args e

171171 (map snd newargs, rename newargs e)

172172

173173addFn :: Name -> LDecl -> State LiftState ()

174−addFn fn d

174+addFn fn d

175175 = do LS n i ds done <- get

176176 put (LS n i ((fn, d) : ds) done)

177177

178178makeFn :: [Name] -> LExp -> State LiftState Name

179−makeFn args exp

179+makeFn args exp

180180 = do fn <- getNextName

181181 let (args', exp') = renameArgs args exp

182182 LS n i ds done <- get

@@ -189,7 +189,7 @@ makeFn args exp

189189 return fn

190190

191191liftAll :: [(Name, LDecl)] -> [(Name, LDecl)]

192−liftAll xs =

192+liftAll xs =

193193 let (LS _ _ decls _) = execState (mapM_ liftDef xs) (LS Nothing 0 [] Map.empty) in

194194 decls

195195

@@ -222,6 +222,7 @@ lift env (LForce e) = do e' <- lift env e

222222lift env (LLet n v e) = do v' <- lift env v

223223 e' <- lift (env ++ [n]) e

224224 return (LLet n v' e')

225+lift env (LLam args (LLam args' e)) = lift env (LLam (args ++ args') e)

225226lift env (LLam args e) = do e' <- lift (env ++ args) e

226227 let usedArgs = nub $ usedIn env e'

227228 fn <- makeFn (usedArgs ++ args) e'

@@ -363,15 +364,15 @@ lsubst n new (LCase t e alts) = let e' = lsubst n new e

363364lsubst n new tm = tm

364365

365366rename :: [(Name, Name)] -> LExp -> LExp

366−rename ns tm@(LV x)

367+rename ns tm@(LV x)

367368 = case lookup x ns of

368369 Just n -> LV n

369370 _ -> tm

370−rename ns (LApp t e args)

371+rename ns (LApp t e args)

371372 = let e' = rename ns e

372373 args' = map (rename ns) args in

373374 LApp t e' args'

374−rename ns (LLazyApp fn args)

375+rename ns (LLazyApp fn args)

375376 = let args' = map (rename ns) args in

376377 LLazyApp fn args'

377378rename ns (LLazyExp e) = LLazyExp (rename ns e)

src/IRTS/LangOpts.hs+46−25

@@ -36,9 +36,10 @@ doInline :: LDefs -> LDecl -> LDecl

3636doInline defs d@(LConstructor _ _ _) = d

3737doInline defs (LFun opts topn args exp)

3838 = let inl = evalState (eval [] initEnv [topn] defs exp)

39− (length args)

39+ (length args)

4040 -- do some case floating, which might arise as a result

41− res = caseFloats 10 inl in

41+ -- then, eta contract

42+ res = eta $ caseFloats 10 inl in

4243 case res of

4344 LLam args' body -> LFun opts topn (map snd initNames ++ args') body

4445 _ -> LFun opts topn (map snd initNames) res

@@ -110,10 +111,10 @@ eval stk env rec defs (LCase ty e alts)

110111 -- If they're all lambdas, bind the lambda at the top

111112 let prefix = getLams (map getRHS alts')

112113 case prefix of

113− [] -> return $ LCase ty e' (replaceInAlts e' alts')

114+ [] -> return $ conOpt $ LCase ty e' (replaceInAlts e' alts')

114115 args -> do alts_red <- mapM (dropArgs args) alts'

115− return $ LLam args

116− (LCase ty e' (replaceInAlts e' alts_red))

116+ return $ LLam args

117+ (conOpt (LCase ty e' (replaceInAlts e' alts_red)))

117118eval stk env rec defs (LOp f es)

118119 = unload stk <$> LOp f <$> mapM (eval [] env rec defs) es

119120eval stk env rec defs (LForeign t s args)

@@ -128,7 +129,7 @@ eval stk env rec defs (LLam args sc)

128129 [] -> eval stk' env' rec defs sc

129130 as -> do ns' <- mapM (\n -> do n' <- nextN

130131 return (n, n')) args'

131− LLam (map snd ns') <$>

132+ unload stk' <$> LLam (map snd ns') <$>

132133 eval [] (map (\ (n, n') -> (n, LV n')) ns' ++ env')

133134 rec defs sc

134135eval stk env rec defs var@(LV n)

@@ -157,6 +158,11 @@ evalAlt stk env rec defs (LConstCase c e)

157158evalAlt stk env rec defs (LDefaultCase e)

158159 = LDefaultCase <$> eval stk env rec defs e

159160

161+apply :: [LExp] -> [(Name, LExp)] -> [Name] -> LDefs -> LExp ->

162+ [Name] -> LExp -> State Int LExp

163+apply stk env rec defs var args body

164+ = eval stk env rec defs (LLam args body)

165+

160166dropArgs :: [Name] -> LAlt -> State Int LAlt

161167dropArgs as (LConCase i n es (LLam args rhs))

162168 = do let old = take (length as) args

@@ -173,6 +179,7 @@ dropArgs as (LDefaultCase (LLam args rhs))

173179

174180caseFloat :: LExp -> LExp

175181caseFloat (LApp tc e es) = LApp tc (caseFloat e) (map caseFloat es)

182+caseFloat (LLazyExp e) = LLazyExp (caseFloat e)

176183caseFloat (LForce e) = LForce (caseFloat e)

177184caseFloat (LCon up i n es) = LCon up i n (map caseFloat es)

178185caseFloat (LOp f es) = LOp f (map caseFloat es)

@@ -180,7 +187,7 @@ caseFloat (LLam ns sc) = LLam ns (caseFloat sc)

180187caseFloat (LLet v val sc) = LLet v (caseFloat val) (caseFloat sc)

181188caseFloat (LCase _ (LCase ct exp alts) alts')

182189 | all conRHS alts || length alts == 1

183− = replaceInCase (LCase ct (caseFloat exp) (map (updateWith alts') alts))

190+ = conOpt $ replaceInCase (LCase ct (caseFloat exp) (map (updateWith alts') alts))

184191 where

185192 conRHS (LConCase _ _ _ (LCon _ _ _ _)) = True

186193 conRHS (LConstCase _ (LCon _ _ _ _)) = True

@@ -194,10 +201,19 @@ caseFloat (LCase _ (LCase ct exp alts) alts')

194201 updateWith alts (LDefaultCase rhs) =

195202 LDefaultCase (caseFloat (conOpt (LCase Shared (caseFloat rhs) alts)))

196203

197− conOpt (LCase ct (LCon _ t n args) alts)

198− = pickAlt n args alts

199− conOpt tm = tm

204+caseFloat (LCase ct exp alts')

205+ = conOpt $ replaceInCase (LCase ct (caseFloat exp) (map cfAlt alts'))

206+ where

207+ cfAlt (LConCase i n es rhs) = LConCase i n es (caseFloat rhs)

208+ cfAlt (LConstCase c rhs) = LConstCase c (caseFloat rhs)

209+ cfAlt (LDefaultCase rhs) = LDefaultCase (caseFloat rhs)

210+caseFloat exp = exp

200211

212+-- Case of constructor

213+conOpt :: LExp -> LExp

214+conOpt (LCase ct (LCon _ t n args) alts)

215+ = pickAlt n args alts

216+ where

201217 pickAlt n args (LConCase i n' es rhs : as) | n == n'

202218 = substAll (zip es args) rhs

203219 pickAlt _ _ (LDefaultCase rhs : as) = rhs

@@ -206,14 +222,7 @@ caseFloat (LCase _ (LCase ct exp alts) alts')

206222

207223 substAll [] rhs = rhs

208224 substAll ((n, tm) : ss) rhs = lsubst n tm (substAll ss rhs)

209−

210−caseFloat (LCase ct exp alts')

211− = replaceInCase (LCase ct (caseFloat exp) (map cfAlt alts'))

212− where

213− cfAlt (LConCase i n es rhs) = LConCase i n es (caseFloat rhs)

214− cfAlt (LConstCase c rhs) = LConstCase c (caseFloat rhs)

215− cfAlt (LDefaultCase rhs) = LDefaultCase (caseFloat rhs)

216−caseFloat exp = exp

225+conOpt tm = tm

217226

218227replaceInCase :: LExp -> LExp

219228replaceInCase (LCase ty e alts)

@@ -225,7 +234,7 @@ replaceInAlts exp alts = dropDups $ concatMap (replaceInAlt exp) alts

225234

226235-- Drop overlapping case (arising from case merging of overlapping

227236-- patterns)

228−dropDups (alt@(LConCase _ i n ns) : alts)

237+dropDups (alt@(LConCase _ i n ns) : alts)

229238 = alt : dropDups (filter (notTag i) alts)

230239 where

231240 notTag i (LConCase _ j n ns) = i /= j

@@ -248,9 +257,9 @@ replaceInAlt exp@(LV var) (LDefaultCase (LCase ty (LV var') alts))

248257replaceInAlt exp a = [a]

249258

250259replaceExp :: LExp -> LExp -> LExp -> LExp

251−replaceExp (LCon _ t n args) new (LCon _ t' n' args')

260+replaceExp (LCon _ t n args) new (LCon _ t' n' args')

252261 | n == n' && args == args' = new

253−replaceExp (LCon _ t n args) new (LApp _ (LV n') args')

262+replaceExp (LCon _ t n args) new (LApp _ (LV n') args')

254263 | n == n' && args == args' = new

255264replaceExp old new tm = tm

256265

@@ -271,7 +280,19 @@ getLamPrefix as (LLam args tm : cs)

271280 | otherwise = getLamPrefix as cs

272281getLamPrefix as (_ : cs) = []

273282

274−apply :: [LExp] -> [(Name, LExp)] -> [Name] -> LDefs -> LExp ->

275− [Name] -> LExp -> State Int LExp

276−apply stk env rec defs var args body

277− = eval stk env rec defs (LLam args body)

283+-- eta contract ('\x -> f x' can just be compiled as 'f' when f is local)

284+eta :: LExp -> LExp

285+eta (LApp tc a es) = LApp tc (eta a) (map eta es)

286+eta (LLazyApp n es) = LLazyApp n (map eta es)

287+eta (LLazyExp e) = LLazyExp (eta e)

288+eta (LForce e) = LForce (eta e)

289+eta (LLet n val sc) = LLet n (eta val) (eta sc)

290+eta (LLam args (LApp tc f args'))

291+ | args' == map LV args = eta f

292+eta (LLam args e) = LLam args (eta e)

293+eta (LProj e i) = LProj (eta e) i

294+eta (LCon a t n es) = LCon a t n (map eta es)

295+eta (LCase ct e alts) = LCase ct (eta e) (map (fmap eta) alts)

296+eta (LOp f es) = LOp f (map eta es)

297+eta tm = tm

298+