7282c0a · Idris-dev · git.iko.soy

git.iko.soy/Idris-dev

Some optimisation of case trees

Case floating, merging, case of constructor (some from Andre Santos' PhD thesis)

commit7282c0ae68b90f78e7f272f24fd6754456cdfcd2 authorEdwin Brady <ecb10@st-andrews.ac.uk> dateFri 15 Dec 2017, 08:29 (9 years ago) parentadc43d8

5 files changed+207−43

compare with master

unifiedsplit

libs/prelude/IO.idr+1−1 src/IRTS/CodegenC.hs+7−1 src/IRTS/JavaScript/PrimOp.hs−2 src/IRTS/Lang.hs+95−33 src/IRTS/LangOpts.hs+104−6

libs/prelude/IO.idr+1−1

@@ -101,7 +101,7 @@ unsafePerformPrimIO : PrimIO a -> a

101101prim_io_pure : a -> PrimIO a

102102prim_io_pure x = Prim__IO x

103103

104−-- %inline

104+-- Don't %inline; the compiler treats it specially

105105io_bind : IO' l a -> (a -> IO' l b) -> IO' l b

106106io_bind (MkIO fn)

107107 = \k => MkIO (\w => prim_io_bind (fn w)

src/IRTS/CodegenC.hs+7−1

@@ -220,8 +220,14 @@ bcc i (PROJECT l loc a) = indent i ++ "PROJECT(vm, " ++ creg l ++ ", " ++ show l

220220 ", " ++ show a ++ ");\n"

221221bcc i (PROJECTINTO r t idx)

222222 = indent i ++ creg r ++ " = GETARG(" ++ creg t ++ ", " ++ show idx ++ ");\n"

223+bcc i (CASE True r [(_, alt)] Nothing)

224+ = indent i ++ showCode i alt

225+ where

226+ showCode :: Int -> [BC] -> String

227+ showCode i bc = "{\n" ++ concatMap (bcc (i + 1)) bc ++

228+ indent i ++ "}\n"

223229bcc i (CASE True r code def)

224− | length code < 4 && length code > 0 = showCase i def code

230+ | length code < 6 && length code > 1 = showCase i def code

225231 where

226232 showCode :: Int -> [BC] -> String

227233 showCode i bc = "{\n" ++ concatMap (bcc (i + 1)) bc ++

src/IRTS/JavaScript/PrimOp.hs−2

@@ -26,8 +26,6 @@ data JsPrimTy = PTBool | PTAny deriving (Eq, Ord)

2626type PrimF = [JsExpr] -> JsExpr

2727type PrimDec = (Bool, JsPrimTy, PrimF) -- the bool indicates if bigint library is used or not

2828

29−deriving instance Ord PrimFn

30−

3129primDB :: Map.Map PrimFn PrimDec

3230primDB =

3331 Map.fromList [

src/IRTS/Lang.hs+95−33

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

1515

1616import Control.Monad.State hiding (lift)

1717import Data.Data (Data)

18−import Data.List

18+import Data.List

19+import qualified Data.Map.Strict as Map

1920import Data.Typeable (Typeable)

2021import GHC.Generics (Generic)

2122

@@ -44,26 +45,26 @@ data LExp = LV Name

4445 | LOp PrimFn [LExp]

4546 | LNothing

4647 | LError String

47− deriving Eq

48+ deriving (Eq, Ord)

4849

4950data FDesc = FCon Name

5051 | FStr String

5152 | FUnknown

5253 | FIO FDesc

5354 | FApp Name [FDesc]

54− deriving (Show, Eq)

55+ deriving (Show, Eq, Ord)

5556

5657data Export = ExportData FDesc -- Exported data descriptor (usually string)

5758 | ExportFun Name -- Idris name

5859 FDesc -- Exported function descriptor

5960 FDesc -- Return type descriptor

6061 [FDesc] -- Argument types

61− deriving (Show, Eq)

62+ deriving (Show, Eq, Ord)

6263

6364data ExportIFace = Export Name -- FFI descriptor

6465 String -- interface file

6566 [Export]

66− deriving (Show, Eq)

67+ deriving (Show, Eq, Ord)

6768

6869-- Primitive operators. Backends are not *required* to implement all

6970-- of these, but should report an error if they are unable

@@ -96,12 +97,12 @@ data PrimFn = LPlus ArithTy | LMinus ArithTy | LTimes ArithTy

9697 | LCrash

9798

9899 | LNoOp

99− deriving (Show, Eq, Generic)

100+ deriving (Show, Eq, Ord, Generic)

100101

101102-- Supported target languages for foreign calls

102103

103104data FCallType = FStatic | FObject | FConstructor

104− deriving (Show, Eq)

105+ deriving (Show, Eq, Ord)

105106

106107data FType = FArith ArithTy

107108 | FFunction

@@ -112,24 +113,24 @@ data FType = FArith ArithTy

112113 | FManagedPtr

113114 | FCData

114115 | FAny

115− deriving (Show, Eq)

116+ deriving (Show, Eq, Ord)

116117

117118-- FIXME: Why not use this for all the IRs now?

118119data LAlt' e = LConCase Int Name [Name] e

119120 | LConstCase Const e

120121 | LDefaultCase e

121− deriving (Show, Eq, Functor, Data, Typeable)

122+ deriving (Show, Eq, Ord, Functor, Data, Typeable)

122123

123124type LAlt = LAlt' LExp

124125

125126data LDecl = LFun [LOpt] Name [Name] LExp -- options, name, arg names, def

126127 | LConstructor Name Int Int -- constructor name, tag, arity

127− deriving (Show, Eq)

128+ deriving (Show, Eq, Ord)

128129

129130type LDefs = Ctxt LDecl

130131

131132data LOpt = Inline | NoInline

132− deriving (Show, Eq)

133+ deriving (Show, Eq, Ord)

133134

134135addTags :: Int -> [(Name, LDecl)] -> (Int, [(Name, LDecl)])

135136addTags i ds = tag i ds []

@@ -140,37 +141,71 @@ addTags i ds = tag i ds []

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

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

142143

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

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

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

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

147+ -- multiple times

148+

149+setBaseName :: Name -> State LiftState ()

150+setBaseName n

151+ = do LS _ i ds done <- get

152+ put (LS (Just n) i ds done)

144153

145154lname (NS n x) i = NS (lname n i) x

146155lname (UN n) i = MN i n

147156lname x i = sMN i (showCG x ++ "_lam")

148157

149−liftAll :: [(Name, LDecl)] -> [(Name, LDecl)]

150−liftAll xs = concatMap (\ (x, d) -> lambdaLift x d) xs

151−

152−lambdaLift :: Name -> LDecl -> [(Name, LDecl)]

153−lambdaLift n (LFun opts _ args e)

154− = let (e', (LS _ _ decls)) = runState (lift args e) (LS n 0 []) in

155− (n, LFun opts n args e') : decls

156−lambdaLift n x = [(n, x)]

157−

158158getNextName :: State LiftState Name

159−getNextName = do LS n i ds <- get

160− put (LS n (i + 1) ds)

161− return (lname n i)

159+getNextName

160+ = do LS mn i ds done <- get

161+ let newn = case mn of

162+ Nothing -> lname (sUN "_") i

163+ Just n -> lname n i

164+ put (LS mn (i + 1) ds done)

165+ return newn

166+

167+renameArgs :: [Name] -> LExp -> ([Name], LExp)

168+renameArgs args e

169+ = let newargNames = map (\i -> sMN i "lift") [0..]

170+ newargs = zip args newargNames in

171+ (map snd newargs, rename newargs e)

162172

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

164−addFn fn d = do LS n i ds <- get

165− put (LS n i ((fn, d) : ds))

174+addFn fn d

175+ = do LS n i ds done <- get

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

177+

178+makeFn :: [Name] -> LExp -> State LiftState Name

179+makeFn args exp

180+ = do fn <- getNextName

181+ let (args', exp') = renameArgs args exp

182+ LS n i ds done <- get

183+ case Map.lookup (args', exp') done of

184+ Just fn -> return fn

185+ Nothing ->

186+ do addFn fn (LFun [Inline] fn args' exp')

187+ LS n i ds done <- get

188+ put (LS n i ds (Map.insert (args', exp') fn done))

189+ return fn

190+

191+liftAll :: [(Name, LDecl)] -> [(Name, LDecl)]

192+liftAll xs =

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

194+ decls

195+

196+liftDef :: (Name, LDecl) -> State LiftState ()

197+liftDef (n, LFun opts _ args e) =

198+ do setBaseName n

199+ e' <- lift args e

200+ addFn n (LFun opts n args e')

201+liftDef (n, x) = addFn n x

166202

167203lift :: [Name] -> LExp -> State LiftState LExp

168204lift env (LV v) = return (LV v) -- Lifting happens before these can exist...

169205lift env (LApp tc (LV n) args) = do args' <- mapM (lift env) args

170206 return (LApp tc (LV n) args')

171207lift env (LApp tc f args) = do f' <- lift env f

172− fn <- getNextName

173− addFn fn (LFun [Inline] fn env f')

208+ fn <- makeFn env f'

174209 args' <- mapM (lift env) args

175210 return (LApp tc (LV fn) (map LV env ++ args'))

176211lift env (LLazyApp n args) = do args' <- mapM (lift env) args

@@ -180,8 +215,7 @@ lift env (LLazyExp (LConst c)) = return (LConst c)

180215-- = lift env (LLazyApp f args)

181216lift env (LLazyExp e) = do e' <- lift env e

182217 let usedArgs = nub $ usedIn env e'

183− fn <- getNextName

184− addFn fn (LFun [NoInline] fn usedArgs e')

218+ fn <- makeFn usedArgs e'

185219 return (LLazyApp fn (map LV usedArgs))

186220lift env (LForce e) = do e' <- lift env e

187221 return (LForce e')

@@ -190,8 +224,7 @@ lift env (LLet n v e) = do v' <- lift env v

190224 return (LLet n v' e')

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

192226 let usedArgs = nub $ usedIn env e'

193− fn <- getNextName

194− addFn fn (LFun [Inline] fn (usedArgs ++ args) e')

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

195228 return (LApp False (LV fn) (map LV usedArgs))

196229lift env (LProj t i) = do t' <- lift env t

197230 return (LProj t' i)

@@ -329,6 +362,35 @@ lsubst n new (LCase t e alts) = let e' = lsubst n new e

329362 LCase t e' alts'

330363lsubst n new tm = tm

331364

365+rename :: [(Name, Name)] -> LExp -> LExp

366+rename ns tm@(LV x)

367+ = case lookup x ns of

368+ Just n -> LV n

369+ _ -> tm

370+rename ns (LApp t e args)

371+ = let e' = rename ns e

372+ args' = map (rename ns) args in

373+ LApp t e' args'

374+rename ns (LLazyApp fn args)

375+ = let args' = map (rename ns) args in

376+ LLazyApp fn args'

377+rename ns (LLazyExp e) = LLazyExp (rename ns e)

378+rename ns (LForce e) = LForce (rename ns e)

379+rename ns (LLet v val sc) = LLet v (rename ns val) (rename ns sc)

380+rename ns (LLam args sc) = LLam args (rename ns sc)

381+rename ns (LProj e i) = LProj (rename ns e) i

382+rename ns (LCon lv t cn args) = let args' = map (rename ns) args in

383+ LCon lv t cn args'

384+rename ns (LOp op args) = let args' = map (rename ns) args in

385+ LOp op args'

386+rename ns (LForeign fd rd args)

387+ = let args' = map (\(d, a) -> (d, rename ns a)) args in

388+ LForeign fd rd args'

389+rename ns (LCase t e alts) = let e' = rename ns e

390+ alts' = map (fmap (rename ns)) alts in

391+ LCase t e' alts'

392+rename ns tm = tm

393+

332394instance Show LExp where

333395 show e = show' [] "" e where

334396 show' env ind (LV n) = show n

@@ -357,7 +419,7 @@ instance Show LExp where

357419 atloc (Just l) = "@" ++ show (LV l) ++ ":"

358420

359421 show' env ind (LCase up e alts)

360− = "case" ++ update ++ show' env ind e ++ " of \n" ++ fmt alts

422+ = "case" ++ update ++ "(" ++ show' env ind e ++ ") of \n" ++ fmt alts

361423 where

362424 update = case up of

363425 Shared -> " "

src/IRTS/LangOpts.hs+104−6

@@ -35,13 +35,21 @@ nextN = do i <- get

3535doInline :: LDefs -> LDecl -> LDecl

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

3737doInline defs (LFun opts topn args exp)

38− = let res = evalState (eval [] initEnv [topn] defs exp)

39− (length args) in

40−-- = let res = evalState (inlineWith [topn] (map (\n -> (n, LV n)) args) exp) 0 in

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

39+ (length args)

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

41+ res = caseFloats 10 inl in

4142 case res of

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

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

4445 where

46+ caseFloats 0 tm = tm

47+ caseFloats n tm

48+ = let res = caseFloat tm in

49+ if res == tm

50+ then res

51+ else caseFloats (n-1) res

52+

4553 initNames = zipWith (\n i -> (n, newn n i)) args [0..]

4654 initEnv = map (\(n, n') -> (n, LV n')) initNames

4755 newn (UN n) i = MN i n

@@ -61,7 +69,11 @@ eval :: [LExp] -> [(Name, LExp)] -> [Name] -> LDefs -> LExp -> State Int LExp

6169eval stk env rec defs (LLazyApp n es)

6270 = unload stk <$> LLazyApp n <$> (mapM (eval [] env rec defs) es)

6371eval stk env rec defs (LForce e)

64− = unload stk <$> LForce <$> eval [] env rec defs e

72+ = do e' <- eval [] env rec defs e

73+ case e' of

74+ LLazyExp forced -> return $ unload stk forced

75+ LLazyApp n es -> return $ unload stk (LApp False (LV n) es)

76+ _ -> return (unload stk (LForce e'))

6577eval stk env rec defs (LLazyExp e)

6678 = unload stk <$> LLazyExp <$> eval [] env rec defs e

6779-- Special case for io_bind, because it needs to keep executing the first

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

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

99111 let prefix = getLams (map getRHS alts')

100112 case prefix of

101− [] -> return $ LCase ty e' alts'

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

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

103− return $ LLam args (LCase ty e' alts_red)

115+ return $ LLam args

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

104117eval stk env rec defs (LOp f es)

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

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

@@ -129,6 +142,8 @@ eval stk env rec defs var@(LV n)

129142 Just (LFun opts _ args body) <- lookupCtxtExact n defs,

130143 Inline `elem` opts ->

131144 apply stk env (n : rec) defs var args body

145+ | Just (LConstructor n t a) <- lookupCtxtExact n defs ->

146+ return (LCon Nothing t n stk)

132147 | otherwise -> return (unload stk var)

133148eval stk env rec defs t = return (unload stk t)

134149

@@ -156,6 +171,89 @@ dropArgs as (LDefaultCase (LLam args rhs))

156171 rhs' <- eval [] (zipWith (\ o n -> (o, LV n)) old as) [] emptyContext rhs

157172 return (LDefaultCase rhs')

158173

174+caseFloat :: LExp -> LExp

175+caseFloat (LApp tc e es) = LApp tc (caseFloat e) (map caseFloat es)

176+caseFloat (LForce e) = LForce (caseFloat e)

177+caseFloat (LCon up i n es) = LCon up i n (map caseFloat es)

178+caseFloat (LOp f es) = LOp f (map caseFloat es)

179+caseFloat (LLam ns sc) = LLam ns (caseFloat sc)

180+caseFloat (LLet v val sc) = LLet v (caseFloat val) (caseFloat sc)

181+caseFloat (LCase _ (LCase ct exp alts) alts')

182+ | all conRHS alts || length alts == 1

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

184+ where

185+ conRHS (LConCase _ _ _ (LCon _ _ _ _)) = True

186+ conRHS (LConstCase _ (LCon _ _ _ _)) = True

187+ conRHS (LDefaultCase (LCon _ _ _ _)) = True

188+ conRHS _ = False

189+

190+ updateWith alts (LConCase i n es rhs) =

191+ LConCase i n es (caseFloat (conOpt (LCase Shared (caseFloat rhs) alts)))

192+ updateWith alts (LConstCase c rhs) =

193+ LConstCase c (caseFloat (conOpt (LCase Shared (caseFloat rhs) alts)))

194+ updateWith alts (LDefaultCase rhs) =

195+ LDefaultCase (caseFloat (conOpt (LCase Shared (caseFloat rhs) alts)))

196+

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

198+ = pickAlt n args alts

199+ conOpt tm = tm

200+

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

202+ = substAll (zip es args) rhs

203+ pickAlt _ _ (LDefaultCase rhs : as) = rhs

204+ pickAlt n args (_ : as) = pickAlt n args as

205+ pickAlt n args [] = error "Can't happen pickAlt - impossible case found"

206+

207+ substAll [] rhs = rhs

208+ 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

217+

218+replaceInCase :: LExp -> LExp

219+replaceInCase (LCase ty e alts)

220+ = LCase ty e (replaceInAlts e alts)

221+replaceInCase exp = exp

222+

223+replaceInAlts :: LExp -> [LAlt] -> [LAlt]

224+replaceInAlts exp alts = dropDups $ concatMap (replaceInAlt exp) alts

225+

226+-- Drop overlapping case (arising from case merging of overlapping

227+-- patterns)

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

229+ = alt : dropDups (filter (notTag i) alts)

230+ where

231+ notTag i (LConCase _ j n ns) = i /= j

232+ notTag _ _ = True

233+dropDups (c : alts) = c : dropDups alts

234+dropDups [] = []

235+

236+

237+replaceInAlt :: LExp -> LAlt -> [LAlt]

238+-- In an alternative, if the case appears on the right hand side, replace

239+-- it with the given expression, to preserve sharing

240+replaceInAlt exp@(LV _) (LConCase i con args rhs)

241+ = [LConCase i con args $

242+ replaceExp (LCon Nothing i con (map LV args)) exp rhs]

243+-- if a default case inspects the same variable as the case it's in,

244+-- remove the inspection and replace with the alternatives

245+-- (i.e. merge the inner case block)

246+replaceInAlt exp@(LV var) (LDefaultCase (LCase ty (LV var') alts))

247+ | var == var' = alts

248+replaceInAlt exp a = [a]

249+

250+replaceExp :: LExp -> LExp -> LExp -> LExp

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

252+ | n == n' && args == args' = new

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

254+ | n == n' && args == args' = new

255+replaceExp old new tm = tm

256+

159257-- dropArgs as (LConstCase c rhs) = LConstCase c (dropRHS as rhs)

160258-- dropArgs as (LDefaultCase rhs) = LDefaultCase (dropRHS as rhs)

161259