0 | module Compiler.ANF

  2 | import Compiler.LambdaLift

  4 | import Core.CompileExpr

  5 | import Core.Context

  7 | import Data.SortedSet

  8 | import Data.Vect

 10 | %default covering

 17 |   public export

 18 |   data AVar : Type where

 19 |        ALocal : Int -> AVar

 20 |        ANull : AVar

 22 |   public export

 23 |   data ANF : Type where

 24 |     AV : FC -> AVar -> ANF

 25 |     AAppName : FC -> (lazy : Maybe LazyReason) -> Name -> List AVar -> ANF

 26 |     AUnderApp : FC -> Name -> (missing : Nat) -> (args : List AVar) -> ANF

 27 |     AApp : FC -> (lazy : Maybe LazyReason) -> (closure : AVar) -> (arg : AVar) -> ANF

 28 |     ALet : FC -> (var : Int) -> ANF -> ANF -> ANF

 29 |     ACon : FC -> Name -> ConInfo -> (tag : Maybe Int) -> List AVar -> ANF

 30 |     AOp : {0 arity : Nat} -> FC -> (lazy : Maybe LazyReason) -> PrimFn arity -> Vect arity AVar -> ANF

 31 |         -- ^ we explicitly bind arity here to silence the warning that it shadows

 32 |         --   existing functions called arity.

 33 |     AExtPrim : FC -> (lazy : Maybe LazyReason) -> Name -> List AVar -> ANF

 34 |     AConCase : FC -> AVar -> List AConAlt -> Maybe ANF -> ANF

 35 |     AConstCase : FC -> AVar -> List AConstAlt -> Maybe ANF -> ANF

 36 |     APrimVal : FC -> Constant -> ANF

 37 |     AErased : FC -> ANF

 38 |     ACrash : FC -> String -> ANF

 40 |   public export

 41 |   data AConAlt : Type where

 42 |        MkAConAlt : Name -> ConInfo -> (tag : Maybe Int) -> (args : List Int) ->

 43 |                    ANF -> AConAlt

 45 |   public export

 46 |   data AConstAlt : Type where

 47 |        MkAConstAlt : Constant -> ANF -> AConstAlt

 49 | public export

 50 | data ANFDef : Type where

 51 |      MkAFun : (args : List Int) -> ANF -> ANFDef

 52 |      MkACon : (tag : Maybe Int) -> (arity : Nat) -> (nt : Maybe Nat) -> ANFDef

 53 |      MkAForeign : (ccs : List String) -> (fargs : List CFType) ->

 54 |                   CFType -> ANFDef

 55 |      MkAError : ANF -> ANFDef

 57 | showLazy : Maybe LazyReason -> String

 58 | showLazy = maybe "" $ (" " ++) . show

 61 |   export

 62 |   Show AVar where

 63 |     show (ALocal i) = "v" ++ show i

 64 |     show ANull = "[__]"

 66 |   export

 67 |   Eq AVar where

 68 |     (ALocal i1) == (ALocal i2) = i1 == i2

 69 |     ANull == ANull = True

 70 |     _ == _ = False

 72 |   export

 73 |   Ord AVar where

 74 |     compare (ALocal i1) (ALocal i2) = compare i1 i2

 75 |     compare (ALocal _) ANull = GT

 76 |     compare ANull (ALocal _) = LT

 77 |     compare ANull ANull = EQ

 79 |   export

 80 |   covering

 81 |   Show ANF where

 82 |     show (AV _ v) = show v

 83 |     show (AAppName fc lazy n args)

 84 |         = show n ++ showLazy lazy ++ "(" ++ showSep ", " (map show args) ++ ")"

 85 |     show (AUnderApp fc n m args)

 86 |         = "<" ++ show n ++ " underapp " ++ show m ++ ">(" ++

 87 |           showSep ", " (map show args) ++ ")"

 88 |     show (AApp fc lazy c arg)

 89 |         = show c ++ showLazy lazy ++ " @ (" ++ show arg ++ ")"

 90 |     show (ALet fc x val sc)

 91 |         = "%let v" ++ show x ++ " = (" ++ show val ++ ") in (" ++ show sc ++ ")"

 92 |     show (ACon fc n _ t args)

 93 |         = "%con " ++ show n ++ "(" ++ showSep ", " (map show args) ++ ")"

 94 |     show (AOp fc lazy op args)

 95 |         = "%op " ++ show op ++ showLazy lazy ++ "(" ++ showSep ", " (toList (map show args)) ++ ")"

 96 |     show (AExtPrim fc lazy p args)

 97 |         = "%extprim " ++ show p ++ showLazy lazy ++ "(" ++ showSep ", " (map show args) ++ ")"

 98 |     show (AConCase fc sc alts def)

 99 |         = "%case " ++ show sc ++ " of { "

100 |              ++ showSep "| " (map show alts) ++ " " ++ show def ++ " }"

101 |     show (AConstCase fc sc alts def)

102 |         = "%case " ++ show sc ++ " of { "

103 |              ++ showSep "| " (map show alts) ++ " " ++ show def ++ " }"

104 |     show (APrimVal _ x) = show x

105 |     show (AErased _) = "___"

106 |     show (ACrash _ x) = "%CRASH(" ++ show x ++ ")"

108 |   export

109 |   covering

110 |   Show AConAlt where

111 |     show (MkAConAlt n _ t args sc)

112 |         = "%conalt " ++ show n ++

113 |              "(" ++ showSep ", " (map showArg args) ++ ") => " ++ show sc

114 |       where

115 |         showArg : Int -> String

116 |         showArg i = "v" ++ show i

118 |   export

119 |   covering

120 |   Show AConstAlt where

121 |     show (MkAConstAlt c sc)

122 |         = "%constalt(" ++ show c ++ ") => " ++ show sc

126 | Show ANFDef where

127 |   show (MkAFun args exp) = show args ++ ": " ++ show exp

128 |   show (MkACon tag arity nt)

129 |       = "Constructor tag " ++ show tag ++ " arity " ++ show arity ++ " newtype by " ++ show nt

130 |   show (MkAForeign ccs args ret)

131 |       = "Foreign call " ++ show ccs ++ " " ++

132 |         show args ++ " -> " ++ show ret

133 |   show (MkAError exp) = "Error: " ++ show exp

135 | AVars : Scope -> Type

136 | AVars = All (\_ => Int)

138 | data Next : Type where

140 | nextVar : {auto v : Ref Next Int} ->

141 |           Core Int

143 |     = do i <- get Next

144 |          put Next (i + 1)

145 |          pure i

147 | lookup : {idx : _} -> (0 p : IsVar x idx vs) -> AVars vs -> Int

148 | lookup First (x :: xs) = x

149 | lookup (Later p) (x :: xs) = lookup p xs

151 | bindArgs : {auto v : Ref Next Int} ->

152 |            List ANF -> Core (List (AVar, Maybe ANF))

153 | bindArgs [] = pure []

154 | bindArgs (AV fc var :: xs)

155 |     = do xs' <- bindArgs xs

156 |          pure $ (var, Nothing) :: xs'

157 | bindArgs (AErased fc :: xs)

158 |     = do xs' <- bindArgs xs

159 |          pure $ (ANull, Nothing) :: xs'

160 | bindArgs (x :: xs)

161 |     = do i <- nextVar

162 |          xs' <- bindArgs xs

163 |          pure $ (ALocal i, Just x) :: xs'

165 | letBind : {auto v : Ref Next Int} ->

166 |           FC -> List ANF -> (List AVar -> ANF) -> Core ANF

167 | letBind fc args f

168 |     = do bargs <- bindArgs args

169 |          pure $ doBind [] bargs

170 |   where

171 |     doBind : List AVar -> List (AVar, Maybe ANF) -> ANF

172 |     doBind vs [] = f (reverse vs)

173 |     doBind vs ((ALocal i, Just t) :: xs)

174 |         = ALet fc i t (doBind (ALocal i :: vs) xs)

175 |     doBind vs ((var, _) :: xs) = doBind (var :: vs) xs

178 | mlet : {auto v : Ref Next Int} ->

179 |        FC -> ANF -> (AVar -> ANF) -> Core ANF

180 | mlet fc (AV _ var) sc = pure $ sc var

181 | mlet fc val sc

182 |     = do i <- nextVar

183 |          pure $ ALet fc i val (sc (ALocal i))

186 |   anfArgs : {auto v : Ref Next Int} ->

187 |             FC -> AVars vars ->

188 |             List (Lifted vars) -> (List AVar -> ANF) -> Core ANF

189 |   anfArgs fc vs args f

190 |       = do args' <- traverse (anf vs) args

191 |            letBind fc args' f

193 |   anf : {auto v : Ref Next Int} ->

194 |         AVars vars -> Lifted vars -> Core ANF

195 |   anf vs (LLocal fc p) = pure $ AV fc (ALocal (lookup p vs))

196 |   anf vs (LAppName fc lazy n args)

197 |       = anfArgs fc vs args (AAppName fc lazy n)

198 |   anf vs (LUnderApp fc n m args)

199 |       = anfArgs fc vs args (AUnderApp fc n m)

200 |   anf vs (LApp fc lazy f a)

201 |       = anfArgs fc vs [f, a] $

202 |                 \case

203 |                   [fvar, avar] => AApp fc lazy fvar avar

204 |                   _ => ACrash fc "Can't happen (AApp)"

205 |   anf vs (LLet fc x val sc)

206 |       = do i <- nextVar

207 |            let vs' = i :: vs

208 |            pure $ ALet fc i !(anf vs val) !(anf vs' sc)

209 |   anf vs (LCon fc n ci t args)

210 |       = anfArgs fc vs args (ACon fc n ci t)

211 |   anf vs (LOp {arity} fc lazy op args)

212 |       = do args' <- traverse (anf vs) (toList args)

213 |            letBind fc args'

214 |                 (\args => case toVect arity args of

215 |                                Nothing => ACrash fc "Can't happen (AOp)"

216 |                                Just argsv => AOp fc lazy op argsv)

217 |   anf vs (LExtPrim fc lazy p args)

218 |       = anfArgs fc vs args (AExtPrim fc lazy p)

219 |   anf vs (LConCase fc scr alts def)

220 |       = do scr' <- anf vs scr

221 |            alts' <- traverse (anfConAlt vs) alts

222 |            def' <- traverseOpt (anf vs) def

223 |            mlet fc scr' (\x => AConCase fc x alts' def')

224 |   anf vs (LConstCase fc scr alts def)

225 |       = do scr' <- anf vs scr

226 |            alts' <- traverse (anfConstAlt vs) alts

227 |            def' <- traverseOpt (anf vs) def

228 |            mlet fc scr' (\x => AConstCase fc x alts' def')

229 |   anf vs (LPrimVal fc c) = pure $ APrimVal fc c

230 |   anf vs (LErased fc) = pure $ AErased fc

231 |   anf vs (LCrash fc err) = pure $ ACrash fc err

233 |   anfConAlt : {auto v : Ref Next Int} ->

234 |               AVars vars -> LiftedConAlt vars -> Core AConAlt

235 |   anfConAlt vs (MkLConAlt n ci t args sc)

236 |       = do (is, vs') <- bindArgs args vs

237 |            pure $ MkAConAlt n ci t is !(anf vs' sc)

238 |     where

239 |       bindArgs : (args : List Name) -> AVars vars' ->

240 |                  Core (List Int, AVars (args ++ vars'))

241 |       bindArgs [] vs = pure ([], vs)

242 |       bindArgs (n :: ns) vs

243 |           = do i <- nextVar

244 |                (is, vs') <- bindArgs ns vs

245 |                pure (i :: is, i :: vs')

247 |   anfConstAlt : {auto v : Ref Next Int} ->

248 |                 AVars vars -> LiftedConstAlt vars -> Core AConstAlt

249 |   anfConstAlt vs (MkLConstAlt c sc)

250 |       = pure $ MkAConstAlt c !(anf vs sc)

253 | toANF : LiftedDef -> Core ANFDef

254 | toANF (MkLFun args scope sc)

255 |     = do v <- newRef Next (the Int 0)

256 |          (iargs, vsNil) <- bindArgs args []

257 |          let vs : AVars args = rewrite sym (appendNilRightNeutral args) in

258 |                                       vsNil

259 |          (iargs', vs) <- bindArgs scope vs

260 |          pure $ MkAFun (iargs ++ reverse iargs') !(anf vs sc)

261 |   where

262 |     bindArgs : {auto v : Ref Next Int} ->

263 |                (args : List Name) -> AVars vars' ->

264 |                Core (List Int, AVars (args ++ vars'))

265 |     bindArgs [] vs = pure ([], vs)

266 |     bindArgs (n :: ns) vs

267 |         = do i <- nextVar

268 |              (is, vs') <- bindArgs ns vs

269 |              pure (i :: is, i :: vs')

270 | toANF (MkLCon t a ns) = pure $ MkACon t a ns

271 | toANF (MkLForeign ccs fargs t) = pure $ MkAForeign ccs fargs t

272 | toANF (MkLError err)

273 |     = do v <- newRef Next (the Int 0)

274 |          pure $ MkAError !(anf [] err)

277 | freeVariables : ANF -> SortedSet AVar

278 | freeVariables (AV _ x) = singleton x

279 | freeVariables (AAppName _ _ n args) = fromList args

280 | freeVariables (AUnderApp _ n _ args) = fromList args

281 | freeVariables (AApp _ _ closure arg) = fromList [closure, arg]

282 | freeVariables (ALet _ var value body) =

283 |     union (freeVariables value) (delete (ALocal var) $ freeVariables body)

284 | freeVariables (ACon _ _ _ _ args) = fromList args

285 | freeVariables (AOp _ _ _ args) = fromList $ toList args

286 | freeVariables (AExtPrim _ _ _ args) = fromList args

287 | freeVariables (AConCase _ sc alts mDef) =

288 |     let altsAnf =

289 |         map (\(MkAConAlt _ _ _ args caseBody) =>

290 |                 difference (freeVariables caseBody) (fromList $ ALocal <$> args)) alts in

291 |     let vars : List (SortedSet AVar) = case mDef of

292 |                 Just anf => freeVariables anf :: altsAnf

293 |                 Nothing => altsAnf in

294 |     insert sc $ concat vars

295 | freeVariables (AConstCase _ sc alts mDef) =

296 |     let altsAnf = map (\(MkAConstAlt _ caseBody) => caseBody) alts in

297 |     let anfs : List ANF = case mDef of

298 |                 Just anf => anf :: altsAnf

299 |                 Nothing => altsAnf in

300 |     insert sc $ foldMap freeVariables anfs

301 | freeVariables _ = empty

304 | usedConstructors : ANF -> SortedSet Name

305 | usedConstructors (AV _ x) = empty

306 | usedConstructors (AAppName _ _ n args) = empty

307 | usedConstructors (AUnderApp _ n _ args) = empty

308 | usedConstructors (AApp _ _ closure arg) = empty

309 | usedConstructors (ALet _ var value body) = union (usedConstructors value) (usedConstructors body)

310 | usedConstructors (ACon _ n _ _ args) = singleton n

311 | usedConstructors (AOp _ _ _ args) = empty

312 | usedConstructors (AExtPrim _ _ _ args) = empty

313 | usedConstructors (AConCase _ sc alts mDef) =

314 |     let altsAnf =

315 |         map (\(MkAConAlt _ _ _ args caseBody) => usedConstructors caseBody) alts in

316 |     let anfs : List (SortedSet Name) = case mDef of

317 |                 Just anf => usedConstructors anf :: altsAnf

318 |                 Nothing => altsAnf in

319 |     concat anfs

320 | usedConstructors (AConstCase _ sc alts mDef) =

321 |     let altsAnf = map (\(MkAConstAlt _ caseBody) => caseBody) alts in

322 |     let anfs : List ANF = case mDef of

323 |                 Just anf => anf :: altsAnf

324 |                 Nothing => altsAnf in

325 |     foldMap usedConstructors anfs

326 | usedConstructors _ = empty