mlton/ssa/local-ref.fun

   1 (* Copyright (C) 1999-2005 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 LocalRef (S: SSA_TRANSFORM_STRUCTS): SSA_TRANSFORM =
  10 struct
  11
  12 open S
  13 open Exp Transfer
  14
  15 structure Prim =
  16    struct
  17       open Prim
  18
  19       val isReff: 'a t -> bool =
  20          fn p =>
  21          case name p of
  22             Name.Ref_ref => true
  23           | _ => false
  24    end
  25
  26 structure FuncLattice = FlatLattice (structure Point = Func)
  27
  28 structure GlobalInfo =
  29   struct
  30     datatype t = T of {isGlobalRef: bool,
  31                        funcUses: FuncLattice.t}
  32
  33     fun layout (T {isGlobalRef, funcUses, ...})
  34       = let open Layout
  35         in record [("isGlobalRef", Bool.layout isGlobalRef),
  36                    ("funcUses", FuncLattice.layout funcUses)]
  37         end
  38
  39     local
  40       fun make f (T r) = f r
  41     in
  42       val isGlobalRef = make #isGlobalRef
  43       val funcUses = make #funcUses
  44     end
  45
  46     fun new isGlobalRef = T {isGlobalRef = isGlobalRef,
  47                              funcUses = FuncLattice.new ()}
  48   end
  49
  50 structure Local =
  51   struct
  52     structure L = TwoPointLattice (val bottom = "local"
  53                                    val top = "non local")
  54     open L
  55     val isLocal = isBottom
  56     val nonLocal = makeTop
  57   end
  58
  59 structure VarInfo =
  60   struct
  61     datatype t = T of {reff: (Label.t * Type.t) option,
  62                        assigns: Label.t list ref,
  63                        derefs: Label.t list ref,
  64                        locall: Local.t,
  65                        threadCopyCurrent: {assign: bool ref,
  66                                            deref: bool ref}}
  67
  68     fun layout (T {reff, assigns, derefs, locall,
  69                    threadCopyCurrent = {assign, deref, ...}, ...})
  70       = let open Layout
  71         in record [("reff", Option.layout (tuple2 (Label.layout, Type.layout)) reff),
  72                    ("assigns", List.layout Label.layout (!assigns)),
  73                    ("derefs", List.layout Label.layout (!derefs)),
  74                    ("locall", Local.layout locall),
  75                    ("threadCopyCurrent", record [("assign", Bool.layout (!assign)),
  76                                                  ("deref", Bool.layout (!deref))])]
  77         end
  78
  79     local
  80       fun make f (T r) = f r
  81       fun make' f = (make f, ! o (make f))
  82     in
  83       val reff = make #reff
  84       val (assigns, _) = make' #assigns
  85       val (derefs, _) = make' #derefs
  86       val locall = make #locall
  87       val threadCopyCurrent = make #threadCopyCurrent
  88     end
  89     val isLocal = Local.isLocal o locall
  90     val nonLocal = Local.nonLocal o locall
  91     local
  92       fun make f = f o threadCopyCurrent
  93       fun make' f = (make f, ! o (make f))
  94     in
  95       val (threadCopyCurrentAssign,threadCopyCurrentAssign') = make' #assign
  96       val (threadCopyCurrentDeref,threadCopyCurrentDeref') = make' #deref
  97     end
  98
  99     fun new reff: t = T {reff = reff,
 100                          assigns = ref [],
 101                          derefs = ref [],
 102                          locall = let
 103                                     val locall = Local.new ()
 104                                     val _ = if isSome reff
 105                                               then ()
 106                                               else Local.nonLocal locall
 107                                   in
 108                                     locall
 109                                   end,
 110                          threadCopyCurrent = {assign = ref false,
 111                                               deref = ref false}}
 112   end
 113
 114 structure LabelInfo =
 115   struct
 116     datatype t = T of {reffs: Var.t list ref,
 117                        assigns: Var.t list ref,
 118                        derefs: Var.t list ref,
 119                        preds: Label.t list ref,
 120                        visited: bool ref}
 121
 122     local
 123       fun make f (T r) = f r
 124       fun make' f = (make f, ! o (make f))
 125     in
 126       val (reffs, reffs') = make' #reffs
 127       val (assigns, assigns') = make' #assigns
 128       val (derefs, derefs') = make' #derefs
 129       val (preds, preds') = make' #preds
 130       val (visited, visited') = make' #visited
 131     end
 132
 133     fun new (): t = T {reffs = ref [],
 134                        assigns = ref [],
 135                        derefs = ref [],
 136                        preds = ref [],
 137                        visited = ref false}
 138   end
 139
 140 structure Multi = Multi (S)
 141
 142 fun transform (program: Program.t): Program.t =
 143    let
 144       val program as Program.T {datatypes, globals, functions, main} =
 145          eliminateDeadBlocks program
 146       (* Compute multi *)
 147       val multi = Control.trace (Control.Detail, "multi") Multi.multi
 148       val {usesThreadsOrConts: bool,
 149            funcIsMultiUsed: Func.t -> bool,
 150            labelDoesThreadCopyCurrent: Label.t -> bool, ...} = multi program
 151       (* Initialize globalInfo *)
 152       val {get = globalInfo: Var.t -> GlobalInfo.t,
 153            set = setGlobalInfo, ...} =
 154          Property.getSetOnce
 155          (Var.plist, Property.initFun (fn _ => GlobalInfo.new false))
 156       val varFuncUses = GlobalInfo.funcUses o globalInfo
 157       val _ =
 158          Vector.foreach
 159          (globals, fn Statement.T {var, exp, ...} =>
 160           Option.app (var, fn var =>
 161                       case exp of
 162                          PrimApp {prim, ...} =>
 163                             if Prim.isReff prim
 164                                then setGlobalInfo (var, GlobalInfo.new true)
 165                             else ()
 166                        | _ => ()))
 167       (* Compute funcUses *)
 168       fun addFunc f x =
 169          let
 170             val gi = globalInfo x
 171          in
 172             if GlobalInfo.isGlobalRef gi
 173                then ignore (FuncLattice.lowerBound (GlobalInfo.funcUses gi, f))
 174             else ()
 175          end
 176       val dummy = Func.newNoname ()
 177       val _ =
 178          Vector.foreach
 179          (globals, fn Statement.T {var, exp, ...} =>
 180           let
 181              fun default () = Exp.foreachVar (exp, addFunc dummy)
 182           in
 183              case exp of
 184                 PrimApp {prim, args, ...} =>
 185                    if Prim.isReff prim
 186                       then
 187                          ignore
 188                          (FuncLattice.<= (varFuncUses (valOf var),
 189                                           varFuncUses (Vector.first args)))
 190                    else default ()
 191               | _ => default ()
 192           end)
 193       val _ =
 194          List.foreach
 195          (functions, fn f =>
 196           let
 197              val {name, blocks, ...} = Function.dest f
 198           in
 199              Vector.foreach
 200              (blocks, fn Block.T {statements, transfer, ...} =>
 201               (Vector.foreach (statements, fn Statement.T {exp, ...} =>
 202                                Exp.foreachVar (exp, addFunc name))
 203                ; Transfer.foreachVar (transfer, addFunc name)))
 204           end)
 205       (* Diagnostics *)
 206       val _ =
 207          Control.diagnostics
 208          (fn display =>
 209           let
 210              open Layout
 211           in
 212              display (str "\n\nGlobals:")
 213              ; (Vector.foreach
 214                 (globals, fn Statement.T {var, ...} =>
 215                  Option.app
 216                  (var, fn x =>
 217                   if GlobalInfo.isGlobalRef (globalInfo x)
 218                      then display (seq [Var.layout x,
 219                                         str ": ",
 220                                         GlobalInfo.layout (globalInfo x)])
 221                   else ())))
 222           end)
 223       (* Localize global refs *)
 224       val {get = funcInfo: Func.t -> {locals: Statement.t list ref}, ...} =
 225          Property.get (Func.plist,
 226                        Property.initFun (fn _ => {locals = ref []}))
 227       val globals =
 228          Vector.keepAllMap
 229          (globals, fn (s as Statement.T {var, ...}) =>
 230           case var of
 231              NONE => SOME s
 232            | SOME x =>
 233                 let
 234                    val GlobalInfo.T {isGlobalRef, funcUses} = globalInfo x
 235                 in
 236                    if not isGlobalRef
 237                       then SOME s
 238                    else
 239                       (case FuncLattice.getPoint funcUses of
 240                           NONE => SOME s
 241                         | SOME f =>
 242                              if funcIsMultiUsed f
 243                                 orelse Func.equals (f, dummy)
 244                                 then SOME s
 245                              else
 246                                 (List.push (#locals (funcInfo f), s)
 247                                  ; NONE))
 248                 end)
 249       (* restore and shrink *)
 250       val restore = restoreFunction {globals = globals}
 251       val shrink = shrinkFunction {globals = globals}
 252       (* varInfo *)
 253       val {get = varInfo: Var.t -> VarInfo.t,
 254            set = setVarInfo, ...}
 255         = Property.getSetOnce
 256           (Var.plist, Property.initFun (fn _ => VarInfo.new NONE))
 257       fun nonLocal x = VarInfo.nonLocal (varInfo x)
 258       fun isLocal x = VarInfo.isLocal (varInfo x)
 259       (* labelInfo *)
 260       val {get = labelInfo: Label.t -> LabelInfo.t,
 261            set = setLabelInfo, ...}
 262         = Property.getSetOnce
 263           (Label.plist, Property.initRaise ("localRef.labelInfo", Label.layout))
 264       fun rewrite (f: Function.t, refs): Function.t =
 265          let
 266             val {args, blocks, mayInline, name, raises, returns, start} =
 267                Function.dest f
 268             (* Diagnostics *)
 269             val _ =
 270                Control.diagnostics
 271                (fn display =>
 272                 let
 273                    open Layout
 274                 in
 275                    display (seq [Func.layout name,
 276                                  str " LocalRefs: ",
 277                                  List.layout
 278                                  (fn x =>
 279                                   seq [Var.layout x,
 280                                        str ": ",
 281                                        VarInfo.layout (varInfo x)])
 282                                  refs])
 283                 end)
 284             (* Rewrite. *)
 285             fun rewriteStatement (s: Statement.t as Statement.T {exp, var, ...})
 286                = let
 287                     datatype z = datatype Prim.Name.t
 288                  in
 289                     case exp
 290                        of PrimApp {prim, args, ...}
 291                           => let
 292                                 fun arg n = Vector.sub (args, n)
 293
 294                                 fun rewriteReffAssign rvar var
 295                                    = let
 296                                         val vi = varInfo rvar
 297                                      in
 298                                         if VarInfo.isLocal vi
 299                                            then Statement.T
 300                                               {var = SOME rvar,
 301                                                ty = #2 (valOf (VarInfo.reff vi)),
 302                                                exp = Var var}
 303                                         else s
 304                                      end
 305                                 fun rewriteReff ()
 306                                    = case var
 307                                    of NONE => s
 308                                  | SOME var => rewriteReffAssign var (arg 0)
 309                                 fun rewriteAssign () = rewriteReffAssign (arg 0) (arg 1)
 310                                 fun rewriteDeref rvar
 311                                    = let
 312                                         val vi = varInfo rvar
 313                                      in
 314                                         if VarInfo.isLocal vi
 315                                            then let
 316                                                 in
 317                                                    Statement.T
 318                                                    {var = var,
 319                                                     ty = #2 (valOf (VarInfo.reff vi)),
 320                                                     exp = Var rvar}
 321                                                 end
 322                                         else s
 323                                      end
 324                                 val rewriteDeref
 325                                    = fn () => rewriteDeref (arg 0)
 326                              in
 327                                 case Prim.name prim
 328                                    of Ref_ref => rewriteReff ()
 329                                  | Ref_assign => rewriteAssign ()
 330                                  | Ref_deref => rewriteDeref ()
 331                                  | _ => s
 332                              end
 333                         | _ => s
 334                  end
 335             fun rewriteBlock (Block.T {label, args, statements, transfer})
 336                = let
 337                     val li = labelInfo label
 338                     (* Don't need to rewrite the statements
 339                      * if this block doesn't mention localizable refs.
 340                      *)
 341                     val statements
 342                        = if List.exists (LabelInfo.reffs' li, isLocal)
 343                        orelse
 344                        List.exists (LabelInfo.assigns' li, isLocal)
 345                        orelse
 346                        List.exists (LabelInfo.derefs' li, isLocal)
 347                             then Vector.map (statements, rewriteStatement)
 348                          else statements
 349                  in
 350                     Block.T {label = label,
 351                              args = args,
 352                              statements = statements,
 353                              transfer = transfer}
 354                  end
 355             val blocks = Vector.map (blocks, rewriteBlock)
 356             val f = Function.new {args = args,
 357                                   blocks = blocks,
 358                                   mayInline = mayInline,
 359                                   name = name,
 360                                   raises = raises,
 361                                   returns = returns,
 362                                   start = start}
 363             val f = restore f
 364             val f = shrink f
 365          in
 366             f
 367          end
 368       val functions =
 369          List.revMap
 370          (functions, fn f =>
 371           let
 372              val {name, ...} = Function.dest f
 373              val {locals, ...} = funcInfo name
 374              val locals = !locals
 375              val f =
 376                 if List.isEmpty locals
 377                    then f
 378                 else
 379                    let
 380                       val {args, blocks, mayInline, name, raises, returns,
 381                            start} = Function.dest f
 382                       val locals = Vector.fromListRev locals
 383                       val localsLabel = Label.newNoname ()
 384                       val localsBlock =
 385                          Block.T {label = localsLabel,
 386                                   args = Vector.new0 (),
 387                                   statements = locals,
 388                                   transfer = Goto {dst = start,
 389                                                    args = Vector.new0 ()}}
 390                       val blocks =
 391                          Vector.concat [Vector.new1 localsBlock, blocks]
 392                    in
 393                       Function.new {args = args,
 394                                     blocks = blocks,
 395                                     mayInline = mayInline,
 396                                     name = name,
 397                                     raises = raises,
 398                                     returns = returns,
 399                                     start = localsLabel}
 400                    end
 401              (* Find all localizable refs. *)
 402              val refs = ref []
 403              fun visitStatement label (Statement.T {var, ty, exp})
 404                = let
 405                    val li = labelInfo label
 406                    fun setReff ()
 407                      = Option.app
 408                        (var, fn var =>
 409                         let
 410                           val vi = VarInfo.new (SOME (label, Type.deRef ty))
 411                           val _ = setVarInfo (var, vi)
 412                         in
 413                           List.push (refs, var) ;
 414                           List.push (LabelInfo.reffs li, var)
 415                         end)
 416                    fun setAssign var
 417                      = (List.push (VarInfo.assigns (varInfo var), label) ;
 418                         List.push (LabelInfo.assigns li, var))
 419                    fun setDeref var
 420                      = (List.push (VarInfo.derefs (varInfo var), label) ;
 421                         List.push (LabelInfo.derefs li, var))
 422                    fun default () = Exp.foreachVar (exp, nonLocal)
 423                    datatype z = datatype Prim.Name.t
 424                  in
 425                    case exp
 426                      of PrimApp {prim, args, ...}
 427                       => let
 428                            fun arg n = Vector.sub (args, n)
 429                          in
 430                            case Prim.name prim
 431                              of Ref_ref => (setReff (); default ())
 432                               | Ref_assign => (setAssign (arg 0);
 433                                                nonLocal (arg 1))
 434                               | Ref_deref => setDeref (arg 0)
 435                               | _ => default ()
 436                          end
 437                       | _ => default ()
 438                  end
 439              fun visitBlock (Block.T {label, statements, transfer, ...})
 440                = let
 441                    val li = LabelInfo.new ()
 442                    val _ = setLabelInfo (label, li)
 443                    val _ = Vector.foreach (statements, visitStatement label)
 444                    val _ = Transfer.foreachVar (transfer, nonLocal)
 445                  in
 446                    if usesThreadsOrConts
 447                      then fn () => Transfer.foreachLabel
 448                                    (transfer, fn l =>
 449                                     List.push (LabelInfo.preds (labelInfo l), label))
 450                      else fn () => ()
 451                  end
 452              val _ = Function.dfs (f, visitBlock)
 453              val refs = List.keepAll (!refs, isLocal)
 454              (* Thread criteria *)
 455              val refs
 456                = if usesThreadsOrConts
 457                    then (List.foreach
 458                          (refs, fn x =>
 459                           let
 460                             val vi = varInfo x
 461                             val def = #1 (valOf (VarInfo.reff vi))
 462                             fun doit (threadCopyCurrent, uses)
 463                               = let
 464                                   val visited = ref []
 465                                   fun doit' l
 466                                     = let
 467                                         val li = labelInfo l
 468                                       in
 469                                         if LabelInfo.visited' li
 470                                           then ()
 471                                           else (List.push (visited, l);
 472                                                 LabelInfo.visited li := true;
 473                                                 if labelDoesThreadCopyCurrent l
 474                                                   then threadCopyCurrent := true
 475                                                   else ();
 476                                                 if Label.equals (def, l)
 477                                                   then ()
 478                                                   else List.foreach
 479                                                        (LabelInfo.preds' li, doit'))
 480                                       end
 481                                 in
 482                                   List.foreach
 483                                   (uses, fn l =>
 484                                    List.foreach
 485                                    (LabelInfo.preds' (labelInfo l), doit')) ;
 486                                   List.foreach
 487                                   (!visited, fn l =>
 488                                    LabelInfo.visited (labelInfo l) := false)
 489                                 end
 490                             val _ = doit (VarInfo.threadCopyCurrentAssign vi,
 491                                           !(VarInfo.assigns vi))
 492                             val _ = doit (VarInfo.threadCopyCurrentDeref vi,
 493                                           !(VarInfo.derefs vi))
 494                           in
 495                             if VarInfo.threadCopyCurrentAssign' vi
 496                                andalso
 497                                VarInfo.threadCopyCurrentDeref' vi
 498                               then VarInfo.nonLocal vi
 499                               else ()
 500                           end);
 501                          List.keepAll (refs, isLocal))
 502                    else refs
 503           in
 504              if 0 < List.length refs
 505                 then rewrite (f, refs)
 506              else
 507                 (Function.clear f
 508                  ; (Control.diagnostics
 509                     (fn display =>
 510                      let
 511                         open Layout
 512                      in
 513                         display (seq [Func.layout name,
 514                                       str " NoLocalRefs"])
 515                      end))
 516                  ; f)
 517           end)
 518       val program = Program.T {datatypes = datatypes,
 519                                globals = globals,
 520                                functions = functions,
 521                                main = main}
 522       val _ = Program.clearTop program
 523     in
 524       program
 525     end
 526 end