mlton/xml/call-count.fun

   1 (* Copyright (C) 1999-2006 Henry Cejtin, Matthew Fluet, Suresh
   2  *    Jagannathan, and Stephen Weeks.
   3  * Copyright (C) 1997-2000 NEC Research Institute.
   4  *
   5  * MLton is released under a BSD-style license.
   6  * See the file MLton-LICENSE for details.
   7  *)
   8
   9 functor CallCount(S: CALL_COUNT_STRUCTS): CALL_COUNT =
  10 struct
  11
  12 open S
  13 open Dec PrimExp
  14
  15 fun instrument(program as Program.T{datatypes, body}, passName: string) =
  16    if !Control.instrument
  17       then
  18          let
  19             datatype kind =
  20                None
  21              | Lam of int (* for curried lambdas, how many arrows remain *)
  22              | Prim
  23
  24             val {get = kind: Var.t -> kind, set} =
  25                Property.new(Var.plist, Property.initConst None)
  26
  27             fun makeLam(l: Lambda.t): kind =
  28                let
  29                   fun loop(l, n) =
  30                      let val {decs, result} = Exp.dest(Lambda.body l)
  31                      in case decs of
  32                         [MonoVal{var, exp = Lambda l, ...}] =>
  33                            if Var.equals(var, VarExp.var result)
  34                               then loop(l, n + 1)
  35                            else n
  36                       | _ => n
  37                      end
  38                in Lam(loop(l, 0))
  39                end
  40
  41             fun inc(name: string) : unit -> Dec.t =
  42                let
  43                   val exp =
  44                      PrimApp
  45                      {prim = Prim.newNullary(concat["MLTON_inc", passName, name]),
  46                       targs = [], args = []}
  47                in fn () =>
  48                   MonoVal{var = Var.newNoname(), ty = Type.unit, exp = exp}
  49                end
  50
  51             val incCount = inc "Unknown"
  52             val incObvious = inc "Known"
  53
  54             val program = Program.T{datatypes = datatypes,
  55                                     body = body}
  56             fun loopExp(e: Exp.t): Exp.t =
  57                let val {decs, result} = Exp.dest e
  58                in Exp.new{decs = loopDecs decs,
  59                           result = result}
  60                end
  61             and loopDecs(ds: Dec.t list): Dec.t list =
  62                case ds of
  63                   [] => []
  64                 | d :: ds =>
  65                    case d of
  66                       MonoVal{var, ty, exp} =>
  67                          let
  68                             fun keep exp =
  69                                MonoVal{var = var, ty = ty, exp = exp}
  70                                :: loopDecs ds
  71                          in case exp of
  72                              App{func, ...} =>
  73                                 let fun rest() = d :: loopDecs ds
  74                                 in case kind(VarExp.var func) of
  75                                    None => incCount() :: rest()
  76                                  | Prim => rest()
  77                                  | Lam n => (if n >= 0
  78                                                 then set(var, Lam(n - 1))
  79                                              else ()
  80                                              ; incObvious() :: rest())
  81                                 end
  82                            | Lambda l =>
  83                                 (set(var,
  84                                      case Exp.decs(Lambda.body l) of
  85                                         [MonoVal{exp = PrimApp _, ...}] => Prim
  86                                       | _ => makeLam l)
  87                                  ; keep(Lambda(loopLambda l)))
  88                            | Case{test, cases, default} =>
  89                                 keep
  90                                 (Case{test = test,
  91                                       cases = List.map(cases, fn (p, e) =>
  92                                                        (p, loopExp e)),
  93                                       default = Option.map loopExp default})
  94                            | Handle{try, catch, handler} =>
  95                                 keep(Handle{try = loopExp try,
  96                                             catch = catch,
  97                                             handler = loopExp handler})
  98                            | _ => d :: loopDecs ds
  99                          end
 100                     | PolyVal{var, tyvars, ty, exp} =>
 101                          PolyVal{var = var, tyvars = tyvars, ty = ty,
 102                                  exp = loopExp exp}
 103                          :: loopDecs ds
 104                     | Fun{tyvars, decs} =>
 105                          (List.foreach(decs, fn {var, lambda, ...} =>
 106                                        set(var, makeLam lambda))
 107                          ; Fun{tyvars = tyvars,
 108                                decs = List.map(decs, fn {var, ty, lambda} =>
 109                                                {var = var, ty = ty,
 110                                                 lambda = loopLambda lambda})}
 111                           :: loopDecs ds)
 112                     | Exception _ => d :: loopDecs ds
 113             and loopLambda(l: Lambda.t): Lambda.t =
 114                let val {arg, argType, body} = Lambda.dest l
 115                in Lambda.new{arg = arg, argType = argType,
 116                              body = loopExp body}
 117                end
 118
 119             val program = Program.T{datatypes = datatypes,
 120                                     body = loopExp body}
 121
 122          in Program.clear program
 123             ; program
 124          end
 125    else program
 126 end