0 | {-
  1 | Optimisations which rewrite constructors and cases to more optimised versions.
  2 | In general optimisations here may change the representation of data types
  3 | so after making changes here you may need to update the ttc version
  4 |
  5 | Nat hack - constructors with ZERO or SUCC ConInfo are replaced with 0 : Integer and (1 +) : Integer -> Integer respectively
  6 |          - cases on ZERO and SUCC are replaced with checking if the integer is 0
  7 |
  8 | Enums - enums (ie data types where all constructors are nullary after erasure) are replaced with a finite integer (Bits8, Bits16 or Bits32)
  9 |
 10 | Unit - case blocks scrutinising UNIT are removed (possibly leaving behind a non-inlinable let binding)
 11 |
 12 | Intrinsic constructors: NIL/CONS, NOTHING/JUST - these are left as constructors,
 13 |     but the name and tag are based on the ConInfo rather than the original name and tag.
 14 |     new `CDef`s must be generated for each of these constructors - see intrinsicCons.
 15 | -}
 16 |
 17 | module Compiler.Opts.Constructor
 18 |
 19 | import Core.CompileExpr
 20 | import Core.Context
 21 | import Data.Vect
 22 |
 23 | export
 24 | data NextMN : Type where
 25 |
 26 | export
 27 | newMN : Ref NextMN Int => String -> Core Name
 28 | newMN base = do
 29 |     x <- get NextMN
 30 |     put NextMN $ x + 1
 31 |     pure $ MN base x
 32 |
 33 |
 34 | -- Rewrite applications of Nat-like constructors and functions to more optimal
 35 | -- versions using Integer
 36 |
 37 | -- None of these should be hard coded, but we'll do it this way until we
 38 | -- have a more general approach to optimising data types!
 39 | -- NOTE: Make sure that names mentioned here are listed in 'natHackNames' in
 40 | -- Common.idr, so that they get compiled, as they won't be spotted by the
 41 | -- usual calls to 'getRefs'.
 42 | data Magic : Type where
 43 |     MagicCRef :
 44 |         Name -> (arity : Nat) -> -- checks
 45 |         (FC -> FC -> forall vars. Vect arity (CExp vars) -> CExp vars) -> --translation
 46 |         Magic
 47 |
 48 | magic : List Magic -> CExp vars -> Maybe (CExp vars)
 49 | magic ms (CLam fc x exp) = CLam fc x <$> magic ms exp
 50 | magic ms e = go ms e where
 51 |
 52 |     fire : Magic -> CExp vars -> Maybe (CExp vars)
 53 |     fire (MagicCRef n arity f) (CApp fc (CRef fc' n') es) = do
 54 |         guard (n == n')
 55 |         map (f fc fc') (toVect arity es)
 56 |     fire _ _ = Nothing
 57 |
 58 |     go : List Magic -> CExp vars -> Maybe (CExp vars)
 59 |     go [] e = Nothing
 60 |     go (m :: ms) e = case fire m e of
 61 |         Nothing => go ms e
 62 |         Just e' => Just e'
 63 |
 64 | magic__integerToNat : FC -> FC -> forall vars. Vect 1 (CExp vars) -> CExp vars
 65 | magic__integerToNat fc fc' [k]
 66 |   = CApp fc (CRef fc' (NS typesNS (UN $ Basic "prim__integerToNat"))) [k]
 67 |
 68 | magic__natMinus : FC -> FC -> forall vars. Vect 2 (CExp vars) -> CExp vars
 69 | magic__natMinus fc fc' [m, n]
 70 |   = magic__integerToNat fc fc'
 71 |     [COp fc (Sub IntegerType) [m, n]]
 72 |
 73 | -- We don't reuse natMinus here because we assume that unsuc will only be called
 74 | -- on S-headed numbers so we do not need the truncating integerToNat call!
 75 | magic__natUnsuc : FC -> FC -> forall vars. Vect 1 (CExp vars) -> CExp vars
 76 | magic__natUnsuc fc fc' [m]
 77 |   = COp fc (Sub IntegerType) [m, CPrimVal fc (BI 1)]
 78 |
 79 | -- TODO: next release remove this and use %builtin pragma
 80 | natHack : List Magic
 81 | natHack =
 82 |     [ MagicCRef (NS typesNS (UN $ Basic "natToInteger")) 1 (\ _, _, [k] => k)
 83 |     , MagicCRef (NS typesNS (UN $ Basic "integerToNat")) 1 magic__integerToNat
 84 |     , MagicCRef (NS typesNS (UN $ Basic "plus")) 2
 85 |          (\ fc, fc', [m,n] => COp fc (Add IntegerType) [m, n])
 86 |     , MagicCRef (NS typesNS (UN $ Basic "mult")) 2
 87 |          (\ fc, fc', [m,n] => COp fc (Mul IntegerType) [m, n])
 88 |     , MagicCRef (NS typesNS (UN $ Basic "minus")) 2 magic__natMinus
 89 |     , MagicCRef (NS typesNS (UN $ Basic "equalNat")) 2
 90 |          (\ fc, fc', [m,n] => COp fc (EQ IntegerType) [m, n])
 91 |     , MagicCRef (NS typesNS (UN $ Basic "compareNat")) 2
 92 |          (\ fc, fc', [m,n] => CApp fc (CRef fc' (NS eqOrdNS (UN $ Basic "compareInteger"))) [m, n])
 93 |     ]
 94 |
 95 | -- get all builtin transformations
 96 | builtinMagic : forall vars. CExp vars -> Maybe (CExp vars)
 97 | builtinMagic = magic natHack
 98 |
 99 | natBranch :  CConAlt vars -> Bool
100 | natBranch (MkConAlt n ZERO _ _ _) = True
101 | natBranch (MkConAlt n SUCC _ _ _) = True
102 | natBranch _ = False
103 |
104 | trySBranch : CExp vars -> CConAlt vars -> Maybe (CExp vars)
105 | trySBranch n (MkConAlt nm SUCC _ [arg] sc)
106 |     = Just (CLet (getFC n) arg YesInline (magic__natUnsuc (getFC n) (getFC n) [n]) sc)
107 | trySBranch _ _ = Nothing
108 |
109 | tryZBranch : CConAlt vars -> Maybe (CExp vars)
110 | tryZBranch (MkConAlt n ZERO _ [] sc) = Just sc
111 | tryZBranch _ = Nothing
112 |
113 | getSBranch : CExp vars -> List (CConAlt vars) -> Maybe (CExp vars)
114 | getSBranch n [] = Nothing
115 | getSBranch n (x :: xs) = trySBranch n x <|> getSBranch n xs
116 |
117 | getZBranch : List (CConAlt vars) -> Maybe (CExp vars)
118 | getZBranch [] = Nothing
119 | getZBranch (x :: xs) = tryZBranch x <|> getZBranch xs
120 |
121 | -- Rewrite case trees on Nat to be case trees on Integer
122 | nat : {auto s : Ref NextMN Int} -> CExp vars -> Core (Maybe (CExp vars))
123 | nat (CCon fc _ ZERO _ []) = pure $ Just $ CPrimVal fc (BI 0)
124 | nat (CCon fc _ SUCC _ [x]) = pure $ Just $ COp fc (Add IntegerType) [CPrimVal fc (BI 1), x]
125 | nat (CConCase fc sc@(CLocal {}) alts def) =
126 |     pure $ if any natBranch alts
127 |             then
128 |                 let defb = fromMaybe (CCrash fc "Nat case not covered") def
129 |                     salt = maybe defb id (getSBranch sc alts)
130 |                     zalt = maybe defb id (getZBranch alts)
131 |                 in Just (CConstCase fc sc [MkConstAlt (BI 0) zalt] (Just salt))
132 |             else Nothing
133 | nat (CConCase fc sc alts def) = do
134 |     x <- newMN "succ"
135 |     Just t <- nat (CConCase fc (CLocal fc First) (map weaken alts) (map weaken def))
136 |         | Nothing => pure Nothing
137 |     pure $ Just $ CLet fc x YesInline sc t
138 | nat _ = pure Nothing
139 |
140 | {-
141 | =========
142 |   Enums
143 | =========
144 | -}
145 |
146 | enumTag : Nat -> Int -> Constant
147 | enumTag k i =
148 |   if      k <= 0xff   then B8 (cast i)
149 |   else if k <= 0xffff then B16 (cast i)
150 |   else                     B32 (cast i)
151 |
152 | enum : CExp vars -> Maybe (CExp vars)
153 | enum (CCon fc _ (ENUM n) (Just tag) []) = Just (CPrimVal fc (enumTag n tag))
154 | enum (CConCase fc sc alts def) = do
155 |     alts' <- traverse toEnum alts
156 |     Just $ CConstCase fc sc alts' def
157 |   where
158 |     toEnum : CConAlt vars -> Maybe (CConstAlt vars)
159 |     toEnum (MkConAlt nm (ENUM n) (Just tag) [] sc)
160 |         = pure $ MkConstAlt (enumTag n tag) sc
161 |     toEnum _ = Nothing
162 | enum t = Nothing
163 |
164 | {-
165 | ========
166 |   Unit
167 | ========
168 | -}
169 |
170 | unitTree : Ref NextMN Int => CExp vars -> Core (Maybe (CExp vars))
171 | unitTree exp@(CConCase fc sc alts def) =
172 |     let [MkConAlt _ UNIT _ [] e] = alts
173 |             | _ => pure Nothing
174 |     in case sc of -- TODO: Check scrutinee has no effect, and skip let binding
175 |         CLocal {} => pure $ Just e
176 |         _ => pure $ Just $ CLet fc !(newMN "_unit") NotInline sc (weaken e)
177 | unitTree t = pure Nothing
178 |
179 | {-
180 | ==========================
181 |   Intrinsic constructors
182 | ==========================
183 | -}
184 |
185 | mkIntrinsicName : String -> Name
186 | mkIntrinsicName x = mkNamespacedName (Just intrinsicNS) (Basic x)
187 |   where
188 |     intrinsicNS : Namespace
189 |     intrinsicNS = mkNamespace "_builtin"
190 |
191 | conInfoNameTag : ConInfo -> Maybe (Name, Int)
192 | conInfoNameTag NIL = Just (mkIntrinsicName "NIL", 0)
193 | conInfoNameTag CONS = Just (mkIntrinsicName "CONS", 1)
194 | conInfoNameTag NOTHING = Just (mkIntrinsicName "NOTHING", 0)
195 | conInfoNameTag JUST = Just (mkIntrinsicName "JUST", 1)
196 | conInfoNameTag _ = Nothing
197 |
198 | export
199 | intrinsicCons : List (Name, FC, CDef)
200 | intrinsicCons = process
201 |     [ (NIL, 0)
202 |     , (CONS, 2)
203 |     , (NOTHING, 0)
204 |     , (JUST, 1)
205 |     ]
206 |   where
207 |     process : List (ConInfo, Nat) -> List (Name, FC, CDef)
208 |     process = mapMaybe $ \(ci, ar) =>
209 |         map (\(n, tag) => (n, (EmptyFC, MkCon (Just tag) ar Nothing))) $ conInfoNameTag ci
210 |
211 | tryIntrinsic : CExp vars -> Maybe (CExp vars)
212 | tryIntrinsic (CCon fc _ ci _ xs) =
213 |     conInfoNameTag ci <&> \(n, tag) => CCon fc n ci (Just tag) xs
214 | tryIntrinsic (CConCase fc e alts def) =
215 |     traverse go alts
216 |         <&> \alts => CConCase fc e alts def
217 |   where
218 |     go : CConAlt vars -> Maybe (CConAlt vars)
219 |     go (MkConAlt _ ci _ as e) =
220 |         conInfoNameTag ci <&> \(n, tag) => MkConAlt n ci (Just tag) as e
221 | tryIntrinsic _ = Nothing
222 |
223 | parameters (try : forall vars. CExp vars -> Core (CExp vars))
224 |     rewriteCExp : CExp vars -> Core (CExp vars)
225 |     -- Recursively rewrite the sub-terms of a CExp
226 |     rewriteSub : CExp vars -> Core (CExp vars)
227 |
228 |     rewriteCConAlt : CConAlt vars -> Core (CConAlt vars)
229 |     rewriteCConstAlt : CConstAlt vars -> Core (CConstAlt vars)
230 |
231 |     rewriteCExp exp = do
232 |         exp' <- rewriteSub exp
233 |         try exp'
234 |
235 |     rewriteSub (CLam fc x e) =
236 |         CLam fc x <$> rewriteCExp e
237 |     rewriteSub (CLet fc x inl e1 e2) =
238 |         CLet fc x inl <$> rewriteCExp e1 <*> rewriteCExp e2
239 |     rewriteSub (CApp fc f xs) =
240 |         CApp fc <$> rewriteCExp f <*> traverse rewriteCExp xs
241 |     rewriteSub (CCon fc n ci t xs) =
242 |         CCon fc n ci t <$> traverse rewriteCExp xs
243 |     rewriteSub (COp fc op xs) =
244 |         COp fc op <$> traverseVect rewriteCExp xs
245 |     rewriteSub (CExtPrim fc n xs) =
246 |         CExtPrim fc n <$> traverse rewriteCExp xs
247 |     rewriteSub (CForce fc lr e) =
248 |         CForce fc lr <$> rewriteCExp e
249 |     rewriteSub (CDelay fc lr e) =
250 |         CDelay fc lr <$> rewriteCExp e
251 |     rewriteSub (CConCase fc x alts def) =
252 |         CConCase fc
253 |             <$> rewriteCExp x
254 |             <*> traverse rewriteCConAlt alts
255 |             <*> traverseOpt rewriteCExp def
256 |     rewriteSub (CConstCase fc x alts def) =
257 |         CConstCase fc
258 |             <$> rewriteCExp x
259 |             <*> traverse rewriteCConstAlt alts
260 |             <*> traverseOpt rewriteCExp def
261 |     rewriteSub e = pure e
262 |
263 |     rewriteCConAlt (MkConAlt n ci t as e) = MkConAlt n ci t as <$> rewriteCExp e
264 |     rewriteCConstAlt (MkConstAlt x e) = MkConstAlt x <$> rewriteCExp e
265 |
266 | sequence : List (forall vars. CExp vars -> Core (Maybe (CExp vars))) -> CExp vars -> Core (CExp vars)
267 | sequence [] e = pure e
268 | sequence (x :: xs) e = do
269 |     e' <- fromMaybe e <$> x e
270 |     sequence xs e'
271 |
272 | export
273 | constructorCExp : Ref NextMN Int => CExp vars -> Core (CExp vars)
274 | constructorCExp = rewriteCExp rw
275 |   where
276 |     rw : forall vars. CExp vars -> Core (CExp vars)
277 |     rw = sequence
278 |         [ pure . builtinMagic
279 |         , nat
280 |         , pure . enum
281 |         , unitTree
282 |         , pure . tryIntrinsic
283 |         ]
284 |
285 |
286 | export
287 | constructorCDef : Ref NextMN Int => CDef -> Core CDef
288 | constructorCDef (MkFun args e) = MkFun args <$> constructorCExp e
289 | constructorCDef def = pure def
290 |