1 (* Copyright (C) 1999-2006, 2008 Henry Cejtin, Matthew Fluet, Suresh
2 * Jagannathan, and Stephen Weeks.
3 * Copyright (C) 1997-2000 NEC Research Institute.
5 * MLton is released under a BSD-style license.
6 * See the file MLton-LICENSE for details.
10 * Duplicate a let bound function at each variable reference
11 * if cost is smaller than threshold.
14 functor Polyvariance (S: XML_TRANSFORM_STRUCTS): XML_TRANSFORM =
18 datatype z = datatype Dec.t
19 datatype z = datatype PrimExp.t
25 fun containsArrow t = containsTycon (t, Tycon.arrow)
30 | SOME (t1, t2) => containsArrow t1 orelse isHigherOrder t2
35 ("Polyvariance.isHigherOrder", layout, Bool.layout)
41 fun lambdaSize (Program.T {body, ...}): Lambda.t -> int =
43 val {get = size: Lambda.t -> int, set, ...} =
44 Property.getSetOnce (Lambda.plist,
45 Property.initRaise ("size", Lambda.layout))
46 fun loopExp (e: Exp.t, n: int): int =
48 (Exp.decs e, n, fn (d, n) =>
50 MonoVal {exp, ...} => loopPrimExp (exp, n + 1)
51 | PolyVal {exp, ...} => loopExp (exp, n + 1)
52 | Fun {decs, ...} => Vector.fold (decs, n, fn ({lambda, ...}, n) =>
53 loopLambda (lambda, n))
54 | Exception _ => n + 1)
55 and loopLambda (l: Lambda.t, n): int =
56 let val m = loopExp (Lambda.body l, 0)
59 and loopPrimExp (e: PrimExp.t, n: int): int =
61 Case {cases, default, ...} =>
69 | SOME (e, _) => loopExp (e, n)),
70 fn (e, n) => loopExp (e, n))
72 | Handle {try, handler, ...} =>
73 loopExp (try, loopExp (handler, n + 1))
74 | Lambda l => loopLambda (l, n + 1)
77 val _ = loopExp (body, 0)
82 fun shouldDuplicate (program as Program.T {body, ...}, hofo, small, product)
85 val costs: (Var.t * int * int * int) list ref = ref []
86 val lambdaSize = lambdaSize program
87 fun isOK (var: Var.t, size: int, numOccurrences: int): bool =
88 let val cost = (numOccurrences - 1) * (size - small)
89 in List.push (costs, (var, size, numOccurrences, cost))
92 type info = {numOccurrences: int ref,
93 shouldDuplicate: bool ref}
94 val {get = varInfo: Var.t -> info option, set = setVarInfo, ...} =
95 Property.getSetOnce (Var.plist, Property.initConst NONE)
96 fun new {lambda = _, ty, var}: unit =
97 if not hofo orelse Type.isHigherOrder ty
98 then setVarInfo (var, SOME {numOccurrences = ref 0,
99 shouldDuplicate = ref false})
101 fun loopExp (e: Exp.t, numDuplicates: int): unit =
103 fun loopVar (x: VarExp.t): unit =
104 case varInfo (VarExp.var x) of
106 | SOME {numOccurrences, ...} =>
107 numOccurrences := !numOccurrences + numDuplicates
108 fun loopVars xs = Vector.foreach (xs, loopVar)
109 val {decs, result} = Exp.dest e
111 fn [] => loopVar result
114 MonoVal {var, ty, exp} =>
117 (new {var = var, ty = ty, lambda = l}
120 val body = Lambda.body l
123 NONE => numDuplicates
124 | SOME {numOccurrences,
126 if isOK (var, lambdaSize l,
128 then (shouldDuplicate := true
131 in loopExp (body, numDuplicates)
136 fn e => loopExp (e, numDuplicates)
140 (loopVar func; loopVar arg)
141 | Case {test, cases, default} =>
143 ; Cases.foreach (cases, loopExp)
145 (default, loopExp o #1)))
146 | ConApp {arg, ...} =>
147 Option.app (arg, loopVar)
149 | Handle {try, handler, ...} =>
150 (loopExp try; loopExp handler)
152 Error.bug "Polyvariance.loopExp.loopDecs: unexpected Lambda"
153 | PrimApp {args, ...} => loopVars args
155 | Raise {exn, ...} => loopVar exn
156 | Select {tuple, ...} => loopVar tuple
157 | Tuple xs => loopVars xs
162 | Fun {decs = lambdas, ...} =>
164 val _ = (Vector.foreach (lambdas, new)
168 (lambdas, [], fn ({var, lambda, ...}, dups) =>
169 let val body = Lambda.body lambda
170 in case varInfo var of
172 (loopExp (body, numDuplicates); dups)
175 size = lambdaSize lambda,
184 (dups, 0, fn ({size, ...}, n) => n + size)
188 fn ({info = {numOccurrences, ...}, ...},
189 n) => n + !numOccurrences)
190 in if isOK (if Vector.isEmpty lambdas
191 then Error.bug "Polyvariance.loopExp.loopDecs: empty lambdas"
193 #var (Vector.first lambdas),
194 size, numOccurrences)
198 info = {shouldDuplicate, ...},
200 (shouldDuplicate := true
201 ; loopExp (body, numOccurrences))))
204 (dups, fn {body, ...} =>
205 loopExp (body, numDuplicates))
208 | _ => Error.bug "Polyvariance.loopExp.loopDecs: strange dec"
211 val _ = loopExp (body, 1)
213 List.insertionSort (l, fn ((_, _, _, c), (_, _, _, c')) => c < c')
218 (sort (!costs), fn (x, size, numOcc, c) =>
219 layout (let open Layout
220 in seq [Var.layout x,
221 str " ", Int.layout size,
222 str " ", Int.layout numOcc,
223 str " ", Int.layout c]
229 | SOME {shouldDuplicate, ...} => !shouldDuplicate
232 fun transform (program as Program.T {datatypes, body, overflow},
237 val shouldDuplicate = shouldDuplicate (program, hofo, small, product)
241 duplicates: Var.t list ref
243 val {get = varInfo: Var.t -> info, set = setVarInfo, ...} =
244 Property.getSet (Var.plist,
245 Property.initRaise ("Polyvariance.info", Var.layout))
246 fun loopVar (x: VarExp.t): VarExp.t =
248 (let val x = VarExp.var x
251 | Dup {duplicates, ...} =>
252 let val x' = Var.new x
253 in List.push (duplicates, x')
257 fun loopVars xs = Vector.map (xs, loopVar)
258 fun bind (x: Var.t): Var.t =
259 let val x' = Var.new x
260 in setVarInfo (x, Replace x')
263 fun bindVarType (x, t) = (bind x, t)
264 fun bindPat (Pat.T {con, targs, arg}) =
267 arg = Option.map (arg, bindVarType)}
268 fun new {lambda = _, ty = _, var}: unit =
269 if shouldDuplicate var
270 then setVarInfo (var, Dup {duplicates = ref []})
271 else ignore (bind var)
272 fun loopExp (e: Exp.t): Exp.t =
274 val {decs, result} = Exp.dest e
276 Exp.make (loopDecs (decs, result))
278 and loopLambda (l: Lambda.t): Lambda.t =
280 val {arg, argType, body, mayInline} = Lambda.dest l
282 Lambda.make {arg = bind arg,
285 mayInline = mayInline}
287 and loopDecs (ds: Dec.t list, result): {decs: Dec.t list,
290 [] => {decs = [], result = loopVar result}
293 MonoVal {var, ty, exp} =>
297 val _ = new {var = var, ty = ty, lambda = l}
298 val {decs, result} = loopDecs (ds, result)
302 MonoVal {var = var, ty = ty,
303 exp = Lambda (loopLambda l)}
305 | Dup {duplicates, ...} =>
307 (!duplicates, decs, fn (var, decs) =>
308 MonoVal {var = var, ty = ty,
309 exp = Lambda (loopLambda l)}
311 in {decs = decs, result = result}
318 App {func = loopVar func,
320 | Case {test, cases, default} =>
322 datatype z = datatype Cases.t
329 (bindPat p, loopExp e)))
336 Case {test = loopVar test,
340 (default, fn (e, r) =>
343 | ConApp {con, targs, arg} =>
346 arg = Option.map (arg, loopVar)}
348 | Handle {try, catch, handler} =>
349 Handle {try = loopExp try,
350 catch = bindVarType catch,
351 handler = loopExp handler}
353 Error.bug "Polyvariance.loopDecs: unexpected Lambda"
354 | PrimApp {prim, targs, args} =>
355 PrimApp {prim = prim,
357 args = loopVars args}
359 | Raise {exn, extend} =>
360 Raise {exn = loopVar exn,
362 | Select {tuple, offset} =>
363 Select {tuple = loopVar tuple,
365 | Tuple xs => Tuple (loopVars xs)
366 | Var x => Var (loopVar x)
368 val {decs, result} = loopDecs (ds, result)
369 in {decs = (MonoVal {var = var, ty = ty, exp = exp}
375 val _ = Vector.foreach (decs, new)
376 val {decs = ds, result} = loopDecs (ds, result)
378 ref [Vector.keepAllMap
379 (decs, fn {var, ty, lambda} =>
382 SOME {var = var, ty = ty,
383 lambda = loopLambda lambda}
387 (decs, fn dec as {var, ...} =>
390 | Dup {duplicates, ...} => SOME (dec, !duplicates))
393 (dups, fn ({var, ...}, duplicates) =>
395 (duplicates, fn var' =>
399 (dups, fn ({var = var'', ...}, _) =>
400 if Var.equals (var, var'')
401 then (setVarInfo (var, Replace var')
408 fn (({ty, lambda, ...}, _), var) =>
410 lambda = loopLambda lambda}))
412 val decs = Vector.concat (!ac)
413 in {decs = Fun {tyvars = Vector.new0 (),
417 | _ => Error.bug "Polyvariance.loopDecs: saw bogus dec"
418 val body = loopExp body
420 Option.map (overflow, fn x =>
423 | _ => Error.bug "Polyvariance.duplicate: duplicating Overflow?")
425 Program.T {datatypes = datatypes,
428 val _ = Program.clear program
435 case !Control.polyvariance of
437 | SOME {hofo, rounds, small, product} =>
445 {display = Control.Layouts Program.layouts,
446 name = "duplicate" ^ (Int.toString (n + 1)),
447 stats = Program.layoutStats,
450 thunk = fn () => shrink (transform (p, hofo, small, product))}