/-
Copyright (c) 2019 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Leonardo de Moura
-/
prelude
import init.lean.compiler.ir.basic
import init.lean.compiler.ir.freevars
import init.control.reader
import init.control.conditional

namespace Lean
namespace IR

/- Remark: in the paper "Counting Immutable Beans" the concepts of
   free and live variables coincide because the paper does *not* consider
   join points. For example, consider the function body `B`
   ```
   let x := ctor_0;
   jmp block_1 x
   ```
   in a context where we have the join point `block_1` defined as
   ```
   block_1 (x : obj) : obj :=
   let z := ctor_0 x y;
   ret z
   ``
   The variable `y` is live in the function body `B` since it occurs in
   `block_1` which is "invoked" by `B`.
-/

namespace IsLive
/-
  We use `State Context` instead of `ReaderT Context Id` because we remove
  non local joint points from `Context` whenever we visit them instead of
  maintaining a set of visited non local join points.

  Remark: we don't need to track local join points because we assume there is
  no variable or join point shadowing in our IR.
-/
abbrev M := State LocalContext

@[inline] def visitVar (w : Index) (x : VarId) : M Bool := pure (HasIndex.visitVar w x)
@[inline] def visitJP (w : Index) (x : JoinPointId) : M Bool := pure (HasIndex.visitJP w x)
@[inline] def visitArg (w : Index) (a : Arg) : M Bool := pure (HasIndex.visitArg w a)
@[inline] def visitArgs (w : Index) (as : Array Arg) : M Bool := pure (HasIndex.visitArgs w as)
@[inline] def visitExpr (w : Index) (e : Expr) : M Bool := pure (HasIndex.visitExpr w e)

partial def visitFnBody (w : Index) : FnBody → M Bool
| (FnBody.vdecl x _ v b)    := visitExpr w v <||> visitFnBody b
| (FnBody.jdecl j ys v b)   := visitFnBody v <||> visitFnBody b
| (FnBody.set x _ y b)      := visitVar w x <||> visitArg w y <||> visitFnBody b
| (FnBody.uset x _ y b)     := visitVar w x <||> visitVar w y <||> visitFnBody b
| (FnBody.sset x _ _ y _ b) := visitVar w x <||> visitVar w y <||> visitFnBody b
| (FnBody.setTag x _ b)     := visitVar w x <||> visitFnBody b
| (FnBody.inc x _ _ b)      := visitVar w x <||> visitFnBody b
| (FnBody.dec x _ _ b)      := visitVar w x <||> visitFnBody b
| (FnBody.del x b)          := visitVar w x <||> visitFnBody b
| (FnBody.mdata _ b)        := visitFnBody b
| (FnBody.jmp j ys)         := visitArgs w ys <||> do {
    ctx ← get;
    match ctx.getJPBody j with
    | some b =>
      -- `j` is not a local join point since we assume we cannot shadow join point declarations.
      -- Instead of marking the join points that we have already been visited, we permanently remove `j` from the context.
      set (ctx.eraseJoinPointDecl j) *> visitFnBody b
    | none   =>
      -- `j` must be a local join point. So do nothing since we have already visite its body.
      pure false
  }
| (FnBody.ret x)            := visitArg w x
| (FnBody.case _ x alts)    := visitVar w x <||> alts.anyM (fun alt => visitFnBody alt.body)
| (FnBody.unreachable)      := pure false

end IsLive

/- Return true if `x` is live in the function body `b` in the context `ctx`.

   Remark: the context only needs to contain all (free) join point declarations.

   Recall that we say that a join point `j` is free in `b` if `b` contains
   `FnBody.jmp j ys` and `j` is not local. -/
def FnBody.hasLiveVar (b : FnBody) (ctx : LocalContext) (x : VarId) : Bool :=
(IsLive.visitFnBody x.idx b).run' ctx

abbrev LiveVarSet   := VarIdSet
abbrev JPLiveVarMap := RBMap JoinPointId LiveVarSet (fun j₁ j₂ => j₁.idx < j₂.idx)

instance LiveVarSet.inhabited : Inhabited LiveVarSet := ⟨{}⟩

def mkLiveVarSet (x : VarId) : LiveVarSet :=
RBTree.empty.insert x

namespace LiveVars

abbrev Collector := LiveVarSet → LiveVarSet

@[inline] private def skip : Collector := fun s => s
@[inline] private def collectVar (x : VarId) : Collector := fun s => s.insert x
private def collectArg : Arg → Collector
| (Arg.var x) := collectVar x
| irrelevant  := skip
@[specialize] private def collectArray {α : Type} (as : Array α) (f : α → Collector) : Collector :=
fun s => as.foldl (fun s a => f a s) s
private def collectArgs (as : Array Arg) : Collector :=
collectArray as collectArg
private def accumulate (s' : LiveVarSet) : Collector :=
fun s => s'.fold (fun s x => s.insert x) s
private def collectJP (m : JPLiveVarMap) (j : JoinPointId) : Collector :=
match m.find j with
| some xs => accumulate xs
| none    => skip -- unreachable for well-formed code
private def bindVar (x : VarId) : Collector :=
fun s => s.erase x
private def bindParams (ps : Array Param) : Collector :=
fun s => ps.foldl (fun s p => s.erase p.x) s

def collectExpr : Expr → Collector
| (Expr.ctor _ ys)       := collectArgs ys
| (Expr.reset _ x)       := collectVar x
| (Expr.reuse x _ _ ys)  := collectVar x ∘ collectArgs ys
| (Expr.proj _ x)        := collectVar x
| (Expr.uproj _ x)       := collectVar x
| (Expr.sproj _ _ x)     := collectVar x
| (Expr.fap _ ys)        := collectArgs ys
| (Expr.pap _ ys)        := collectArgs ys
| (Expr.ap x ys)         := collectVar x ∘ collectArgs ys
| (Expr.box _ x)         := collectVar x
| (Expr.unbox x)         := collectVar x
| (Expr.lit v)           := skip
| (Expr.isShared x)      := collectVar x
| (Expr.isTaggedPtr x)   := collectVar x

partial def collectFnBody : FnBody → JPLiveVarMap → Collector
| (FnBody.vdecl x _ v b)    m := collectExpr v ∘ collectFnBody b m ∘ bindVar x
| (FnBody.jdecl j ys v b)   m :=
  let jLiveVars := (collectFnBody v m ∘ bindParams ys) {};
  let m         := m.insert j jLiveVars;
  collectFnBody b m
| (FnBody.set x _ y b)      m := collectVar x ∘ collectArg y ∘ collectFnBody b m
| (FnBody.setTag x _ b)     m := collectVar x ∘ collectFnBody b m
| (FnBody.uset x _ y b)     m := collectVar x ∘ collectVar y ∘ collectFnBody b m
| (FnBody.sset x _ _ y _ b) m := collectVar x ∘ collectVar y ∘ collectFnBody b m
| (FnBody.inc x _ _ b)      m := collectVar x ∘ collectFnBody b m
| (FnBody.dec x _ _ b)      m := collectVar x ∘ collectFnBody b m
| (FnBody.del x b)          m := collectVar x ∘ collectFnBody b m
| (FnBody.mdata _ b)        m := collectFnBody b m
| (FnBody.ret x)            m := collectArg x
| (FnBody.case _ x alts)    m := collectVar x ∘ collectArray alts (fun alt => collectFnBody alt.body m)
| (FnBody.unreachable)      m := skip
| (FnBody.jmp j xs)         m := collectJP m j ∘ collectArgs xs

def updateJPLiveVarMap (j : JoinPointId) (ys : Array Param) (v : FnBody) (m : JPLiveVarMap) : JPLiveVarMap :=
let jLiveVars := (collectFnBody v m ∘ bindParams ys) {};
m.insert j jLiveVars

end LiveVars

def updateLiveVars (e : Expr) (v : LiveVarSet) : LiveVarSet :=
LiveVars.collectExpr e v

def collectLiveVars (b : FnBody) (m : JPLiveVarMap) (v : LiveVarSet := {}) : LiveVarSet :=
LiveVars.collectFnBody b m v

export LiveVars (updateJPLiveVarMap)

end IR
end Lean