From 5ac43d4515343832b16ab00a41953cd9a161d1da Mon Sep 17 00:00:00 2001
From: KDr2 <zhuo.dev@gmail.com>
Date: Sat, 5 Mar 2022 07:50:25 +0800
Subject: [PATCH 01/15] opt: remove box

---
 src/tapedfunction.jl | 121 +++++++++++++++++++------------------------
 1 file changed, 52 insertions(+), 69 deletions(-)

diff --git a/src/tapedfunction.jl b/src/tapedfunction.jl
index becb1e2e..bef60258 100644
--- a/src/tapedfunction.jl
+++ b/src/tapedfunction.jl
@@ -1,13 +1,3 @@
-mutable struct Box{T}
-    id::Symbol
-    val::T
-end
-
-## methods for Box
-Box(x) = Box(gensym(), x)
-Box{T}(x) where {T} = Box{T}(gensym(), x)
-Base.show(io::IO, box::Box) = print(io, "Box(", box.id, ")")
-
 ## Instruction and TapedFunction
 
 abstract type AbstractInstruction end
@@ -41,7 +31,7 @@ mutable struct TapedFunction{F}
     ir::Core.CodeInfo
     tape::RawTape
     counter::Int
-    bindings::Dict{Symbol, Box{<:Any}}
+    bindings::Dict{Symbol, Any}
     retval::Symbol
 
     function TapedFunction(f::F, args...; cache=false) where {F}
@@ -56,12 +46,10 @@ mutable struct TapedFunction{F}
         end
 
         ir = CodeInfoTools.code_inferred(f, args_type...)
-        tf = new{F}() # leave some fields to be undef
-        tf.func, tf.arity, tf.ir = f, length(args), ir
-        tf.tape = RawTape()
-        tf.counter = 1
+        tape = RawTape()
+        bindings = translate!(tape, ir)
 
-        translate!(tf, ir)
+        tf = new{F}(f, length(args), ir, tape, 1, bindings, :none)
         TRCache[cache_key] = tf # set cache
         return tf
     end
@@ -75,9 +63,6 @@ end
 const TRCache = LRU{Tuple, TapedFunction}(maxsize=10)
 
 val(x) = x
-val(x::Box) = x.val
-val(x::Box{GlobalRef}) = val(x.val)
-val(x::Box{QuoteNode}) = val(x.val)
 val(x::GlobalRef) = getproperty(x.mod, x.name)
 val(x::QuoteNode) = eval(x)
 val(x::TapedFunction) = x.func
@@ -86,10 +71,10 @@ result(t::TapedFunction) = val(t.bindings[t.retval])
 function (tf::TapedFunction)(args...; callback=nothing)
     # set args
     if tf.counter <= 1
-        haskey(tf.bindings, :_1) && (tf.bindings[:_1].val = tf.func)
+        haskey(tf.bindings, :_1) && _update_var!(tf, :_1, tf.func)
         for i in 1:length(args)
             slot = Symbol("_", i + 1)
-            haskey(tf.bindings, slot) && (tf.bindings[slot].val = args[i])
+            haskey(tf.bindings, slot) && _update_var!(tf, slot, args[i])
         end
     end
 
@@ -148,6 +133,7 @@ end
 
 _lookup(tf::TapedFunction, v) = v
 _lookup(tf::TapedFunction, v::Symbol) = tf.bindings[v]
+_update_var!(tf::TapedFunction, v::Symbol, c) = tf.bindings[v] = c
 
 function (instr::Instruction{F})(tf::TapedFunction) where F
     # catch run-time exceptions / errors.
@@ -155,8 +141,7 @@ function (instr::Instruction{F})(tf::TapedFunction) where F
         func = val(_lookup(tf, instr.func))
         inputs = map(x -> val(_lookup(tf, x)), instr.input)
         output = func(inputs...)
-        output_box = _lookup(tf, instr.output)
-        output_box.val = output
+        _update_var!(tf, instr.output, output)
         tf.counter += 1
     catch e
         println("counter=", tf.counter)
@@ -201,60 +186,60 @@ end
 
 ## Translation: CodeInfo -> Tape
 
-function var_boxer(var, boxes::Dict{Symbol, Box{<:Any}}) # for literal constants
-    box = Box(var)
-    boxes[box.id] = box
-    return box.id
+function bind_var!(var, bindings::Dict{Symbol, Any}) # for literal constants
+    id = gensym()
+    bindings[id] = var
+    return id
+end
+bind_var!(var::Core.SSAValue, bindings::Dict{Symbol, Any}) = bind_var!(Symbol(var.id), bindings)
+bind_var!(var::Core.TypedSlot, bindings::Dict{Symbol, Any}) =
+    bind_var!(Symbol(:_, var.id), bindings)
+bind_var!(var::Core.SlotNumber, bindings::Dict{Symbol, Any}) =
+    bind_var!(Symbol(:_, var.id), bindings)
+function bind_var!(var::Symbol, bindings::Dict{Symbol, Any})
+    get!(bindings, var, nothing)
+    return var
 end
-var_boxer(var::Core.SSAValue, boxes::Dict{Symbol, Box{<:Any}}) = var_boxer(Symbol(var.id), boxes)
-var_boxer(var::Core.TypedSlot, boxes::Dict{Symbol, Box{<:Any}}) =
-    var_boxer(Symbol(:_, var.id), boxes)
-var_boxer(var::Core.SlotNumber, boxes::Dict{Symbol, Box{<:Any}}) =
-    var_boxer(Symbol(:_, var.id), boxes)
-var_boxer(var::Symbol, boxes::Dict{Symbol, Box{<:Any}}) =
-    get!(boxes, var, Box{Any}(var, nothing)).id
-
-function translate!(tf::TapedFunction, ir::Core.CodeInfo)
-    tape = tf.tape
-    boxes = Dict{Symbol, Box{<:Any}}()
+
+function translate!(tape::RawTape, ir::Core.CodeInfo)
+    bindings = Dict{Symbol, Any}()
 
     for (idx, line) in enumerate(ir.code)
         isa(line, Core.Const) && (line = line.val) # unbox Core.Const
-        ins = translate!!(Core.SSAValue(idx), line, boxes, ir)
+        ins = translate!!(Core.SSAValue(idx), line, bindings, ir)
         push!(tape, ins)
     end
-
-    tf.bindings = boxes
+    return bindings
 end
 
 const IRVar = Union{Core.SSAValue, Core.SlotNumber}
 
 function translate!!(var::IRVar, line::Core.NewvarNode,
-                     boxes::Dict{Symbol, Box{<:Any}}, @nospecialize(ir))
+                     bindings::Dict{Symbol, Any}, @nospecialize(ir))
     # use a noop to ensure the 1-to-1 mapping from ir.code to instructions
     # on tape. see GotoInstruction.dest.
     return GotoInstruction(:_true, 0)
 end
 
 function translate!!(var::IRVar, line::GlobalRef,
-                     boxes::Dict{Symbol, Box{<:Any}}, @nospecialize(ir))
-    return Instruction(() -> val(line), (), var_boxer(var, boxes))
+                     bindings::Dict{Symbol, Any}, @nospecialize(ir))
+    return Instruction(() -> val(line), (), bind_var!(var, bindings))
 end
 
 function translate!!(var::IRVar, line::Core.SlotNumber,
-                     boxes::Dict{Symbol, Box{<:Any}}, @nospecialize(ir))
-    return Instruction(identity, (var_boxer(line, boxes),), var_boxer(var, boxes))
+                     bindings::Dict{Symbol, Any}, @nospecialize(ir))
+    return Instruction(identity, (bind_var!(line, bindings),), bind_var!(var, bindings))
 end
 
 function translate!!(var::IRVar, line::Core.TypedSlot,
-                     boxes::Dict{Symbol, Box{<:Any}}, @nospecialize(ir))
-    input_box = var_boxer(Core.SlotNumber(line.id), boxes)
-    return Instruction(identity, (input_box,), var_boxer(var, boxes))
+                     bindings::Dict{Symbol, Any}, @nospecialize(ir))
+    input_box = bind_var!(Core.SlotNumber(line.id), bindings)
+    return Instruction(identity, (input_box,), bind_var!(var, bindings))
 end
 
 function translate!!(var::IRVar, line::Core.GotoIfNot,
-                     boxes::Dict{Symbol, Box{<:Any}}, @nospecialize(ir))
-    _cond = var_boxer(line.cond, boxes)
+                     bindings::Dict{Symbol, Any}, @nospecialize(ir))
+    _cond = bind_var!(line.cond, bindings)
     cond = if isa(_cond, Bool)
         _cond ? :_true : :_false
     else
@@ -264,24 +249,24 @@ function translate!!(var::IRVar, line::Core.GotoIfNot,
 end
 
 function translate!!(var::IRVar, line::Core.GotoNode,
-                     boxes::Dict{Symbol, Box{<:Any}}, @nospecialize(ir))
+                     bindings::Dict{Symbol, Any}, @nospecialize(ir))
     return GotoInstruction(:_false, line.label)
 end
 
 function translate!!(var::IRVar, line::Core.ReturnNode,
-                     boxes::Dict{Symbol, Box{<:Any}}, @nospecialize(ir))
-    return ReturnInstruction(var_boxer(line.val, boxes))
+                     bindings::Dict{Symbol, Any}, @nospecialize(ir))
+    return ReturnInstruction(bind_var!(line.val, bindings))
 end
 
 function translate!!(var::IRVar, line::Expr,
-                     boxes::Dict{Symbol, Box{<:Any}}, ir::Core.CodeInfo)
+                     bindings::Dict{Symbol, Any}, ir::Core.CodeInfo)
     head = line.head
-    _box_fn = (x) -> var_boxer(x, boxes)
+    _bind_fn = (x) -> bind_var!(x, bindings)
     if head === :new
-        args = map(_box_fn, line.args)
-        return Instruction(__new__, args |> Tuple, var_boxer(var, boxes))
+        args = map(_bind_fn, line.args)
+        return Instruction(__new__, args |> Tuple, bind_var!(var, bindings))
     elseif head === :call
-        args = map(_box_fn, line.args)
+        args = map(_bind_fn, line.args)
         # args[1] is the function
         func = line.args[1]
         if Meta.isexpr(func, :static_parameter) # func is a type parameter
@@ -289,15 +274,15 @@ function translate!!(var::IRVar, line::Expr,
         else # isa(func, GlobalRef) or a var?
             func = args[1] # a var(box)
         end
-        return Instruction(func, args[2:end] |> Tuple, var_boxer(var, boxes))
+        return Instruction(func, args[2:end] |> Tuple, bind_var!(var, bindings))
     elseif head === :(=)
         # line.args[1] (the left hand side) is a SlotNumber, and it should be the output
         lhs = line.args[1]
         rhs = line.args[2] # the right hand side, maybe a Expr, or a var, or ...
         if Meta.isexpr(rhs, (:new, :call))
-            return translate!!(lhs, rhs, boxes, ir)
+            return translate!!(lhs, rhs, bindings, ir)
         else # rhs is a single value
-            return Instruction(identity, (_box_fn(rhs),), _box_fn(lhs))
+            return Instruction(identity, (_bind_fn(rhs),), _bind_fn(lhs))
         end
     else
         @error "Unknown Expression: " typeof(var) var typeof(line) line
@@ -305,12 +290,12 @@ function translate!!(var::IRVar, line::Expr,
     end
 end
 
-function translate!!(var, line, boxes, ir)
+function translate!!(var, line, bindings, ir)
     @error "Unknown IR code: " typeof(var) var typeof(line) line
     throw(ErrorException("Unknown IR code"))
 end
 
-## copy Box, TapedFunction
+## copy Bindings, TapedFunction
 
 """
     tape_copy(x)
@@ -329,12 +314,10 @@ tape_copy(x::Core.Box) = Core.Box(tape_copy(x.contents))
 # tape_copy(x::Array) = deepcopy(x)
 # tape_copy(x::Dict) = deepcopy(x)
 
-Base.copy(box::Box{T}) where T =  Box{T}(box.id, tape_copy(box.val))
-
-function copy_bindings(old::Dict{Symbol, Box{<:Any}})
-    newb = Dict{Symbol, Box{<:Any}}()
+function copy_bindings(old::Dict{Symbol, Any})
+    newb = Dict{Symbol, Any}()
     for (k, v) in old
-        newb[k] = copy(v)
+        newb[k] = tape_copy(v)
     end
     return newb
 end

From cce6cb30d3d3ab7ff0462b516e50860a71d0c559 Mon Sep 17 00:00:00 2001
From: KDr2 <zhuo.dev@gmail.com>
Date: Sat, 5 Mar 2022 08:14:52 +0800
Subject: [PATCH 02/15] use FunctionWrapper as Instruction

---
 Project.toml         |  1 +
 src/Libtask.jl       |  1 +
 src/tapedfunction.jl | 33 ++++++++++-----------------------
 test/tf.jl           |  4 ++--
 4 files changed, 14 insertions(+), 25 deletions(-)

diff --git a/Project.toml b/Project.toml
index 2a88a031..9b133695 100644
--- a/Project.toml
+++ b/Project.toml
@@ -7,6 +7,7 @@ version = "0.7.0"
 
 [deps]
 CodeInfoTools = "bc773b8a-8374-437a-b9f2-0e9785855863"
+FunctionWrappers = "069b7b12-0de2-55c6-9aab-29f3d0a68a2e"
 LRUCache = "8ac3fa9e-de4c-5943-b1dc-09c6b5f20637"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
diff --git a/src/Libtask.jl b/src/Libtask.jl
index fcbb47ea..bea169d0 100644
--- a/src/Libtask.jl
+++ b/src/Libtask.jl
@@ -1,6 +1,7 @@
 module Libtask
 
 using CodeInfoTools
+using FunctionWrappers: FunctionWrapper
 using LRUCache
 
 export TapedTask, consume, produce
diff --git a/src/tapedfunction.jl b/src/tapedfunction.jl
index bef60258..39465ee8 100644
--- a/src/tapedfunction.jl
+++ b/src/tapedfunction.jl
@@ -1,7 +1,6 @@
 ## Instruction and TapedFunction
 
 abstract type AbstractInstruction end
-const RawTape = Vector{AbstractInstruction}
 
 """
     Instruction
@@ -29,7 +28,7 @@ mutable struct TapedFunction{F}
     func::F # maybe a function, a constructor, or a callable object
     arity::Int
     ir::Core.CodeInfo
-    tape::RawTape
+    tape::Vector{FunctionWrapper{Nothing, Tuple{TapedFunction}}}
     counter::Int
     bindings::Dict{Symbol, Any}
     retval::Symbol
@@ -46,7 +45,7 @@ mutable struct TapedFunction{F}
         end
 
         ir = CodeInfoTools.code_inferred(f, args_type...)
-        tape = RawTape()
+        tape = Vector{FunctionWrapper{Nothing, Tuple{TapedFunction}}}()
         bindings = translate!(tape, ir)
 
         tf = new{F}(f, length(args), ir, tape, 1, bindings, :none)
@@ -83,12 +82,15 @@ function (tf::TapedFunction)(args...; callback=nothing)
         ins = tf.tape[tf.counter]
         ins(tf)
         callback !== nothing && callback()
-        isa(ins, ReturnInstruction) && break
+        tf.retval !== :none && break
     end
     return result(tf)
 end
 
 function Base.show(io::IO, tf::TapedFunction)
+    # we use an extra IOBuffer to collect all the data and then
+    # output it once to avoid output interrupt during task context
+    # switching
     buf = IOBuffer()
     println(buf, "TapedFunction:")
     println(buf, "* .func => $(tf.func)")
@@ -103,22 +105,6 @@ function Base.show(io::IO, tf::TapedFunction)
     print(io, String(take!(buf)))
 end
 
-function Base.show(io::IO, rtape::RawTape)
-    # we use an extra IOBuffer to collect all the data and then
-    # output it once to avoid output interrupt during task context
-    # switching
-    buf = IOBuffer()
-    print(buf, length(rtape), "-element RawTape")
-    isempty(rtape) || println(buf, ":")
-    i = 1
-    for instr in rtape
-        print(buf, "\t", i, " => ")
-        show(buf, instr)
-        i += 1
-    end
-    print(io, String(take!(buf)))
-end
-
 ## methods for Instruction
 Base.show(io::IO, instr::AbstractInstruction) = println(io, "A ", typeof(instr))
 
@@ -191,7 +177,8 @@ function bind_var!(var, bindings::Dict{Symbol, Any}) # for literal constants
     bindings[id] = var
     return id
 end
-bind_var!(var::Core.SSAValue, bindings::Dict{Symbol, Any}) = bind_var!(Symbol(var.id), bindings)
+bind_var!(var::Core.SSAValue, bindings::Dict{Symbol, Any}) =
+    bind_var!(Symbol(var.id), bindings)
 bind_var!(var::Core.TypedSlot, bindings::Dict{Symbol, Any}) =
     bind_var!(Symbol(:_, var.id), bindings)
 bind_var!(var::Core.SlotNumber, bindings::Dict{Symbol, Any}) =
@@ -201,13 +188,13 @@ function bind_var!(var::Symbol, bindings::Dict{Symbol, Any})
     return var
 end
 
-function translate!(tape::RawTape, ir::Core.CodeInfo)
+function translate!(tape::Vector{FunctionWrapper{Nothing, Tuple{TapedFunction}}}, ir::Core.CodeInfo)
     bindings = Dict{Symbol, Any}()
 
     for (idx, line) in enumerate(ir.code)
         isa(line, Core.Const) && (line = line.val) # unbox Core.Const
         ins = translate!!(Core.SSAValue(idx), line, bindings, ir)
-        push!(tape, ins)
+        push!(tape, FunctionWrapper{Nothing, Tuple{TapedFunction}}(ins))
     end
     return bindings
 end
diff --git a/test/tf.jl b/test/tf.jl
index b4af5dab..3192460b 100644
--- a/test/tf.jl
+++ b/test/tf.jl
@@ -11,7 +11,7 @@ using Libtask
         tf = Libtask.TapedFunction(S, 1, 2)
         s1 = tf(1, 2)
         @test s1.i == 3
-        newins = findall(x -> isa(x, Libtask.Instruction{typeof(Libtask.__new__)}), tf.tape)
-        @test length(newins) == 1
+        # newins = findall(x -> isa(x, Libtask.Instruction{typeof(Libtask.__new__)}), tf.tape)
+        # @test length(newins) == 1
     end
 end

From f1b7ac3e5a6202969400eaae5f75be64c3cd9f25 Mon Sep 17 00:00:00 2001
From: KDr2 <zhuo.dev@gmail.com>
Date: Sat, 5 Mar 2022 09:23:41 +0800
Subject: [PATCH 03/15] another approach to add type info in

---
 src/tapedfunction.jl | 99 ++++++++++++++++++++++++--------------------
 1 file changed, 55 insertions(+), 44 deletions(-)

diff --git a/src/tapedfunction.jl b/src/tapedfunction.jl
index 39465ee8..69128720 100644
--- a/src/tapedfunction.jl
+++ b/src/tapedfunction.jl
@@ -2,26 +2,30 @@
 
 abstract type AbstractInstruction end
 
+struct TypedVar{T}
+    id::Symbol
+end
+
 """
     Instruction
 
 An `Instruction` stands for a function call
 """
-struct Instruction{F, N} <: AbstractInstruction
+struct Instruction{F, N, TO} <: AbstractInstruction
     func::F
-    input::NTuple{N, Symbol}
-    output::Symbol
+    input::NTuple{N, TypedVar{<:Any}}
+    output::TypedVar{TO}
 end
 
-struct GotoInstruction <: AbstractInstruction
-    condition::Symbol
+struct GotoInstruction{T} <: AbstractInstruction
+    condition::TypedVar{T}
     # we enusre a 1-to-1 mapping between ir.code and instruction
     # so here we can use the index directly.
     dest::Int
 end
 
-struct ReturnInstruction <: AbstractInstruction
-    arg::Symbol
+struct ReturnInstruction{T} <: AbstractInstruction
+    arg::TypedVar{T}
 end
 
 mutable struct TapedFunction{F}
@@ -31,7 +35,7 @@ mutable struct TapedFunction{F}
     tape::Vector{FunctionWrapper{Nothing, Tuple{TapedFunction}}}
     counter::Int
     bindings::Dict{Symbol, Any}
-    retval::Symbol
+    retval::TypedVar
 
     function TapedFunction(f::F, args...; cache=false) where {F}
         args_type = _accurate_typeof.(args)
@@ -48,14 +52,14 @@ mutable struct TapedFunction{F}
         tape = Vector{FunctionWrapper{Nothing, Tuple{TapedFunction}}}()
         bindings = translate!(tape, ir)
 
-        tf = new{F}(f, length(args), ir, tape, 1, bindings, :none)
+        tf = new{F}(f, length(args), ir, tape, 1, bindings, TypedVar{Any}(:none))
         TRCache[cache_key] = tf # set cache
         return tf
     end
 
     function TapedFunction(tf::TapedFunction{F}) where {F}
         new{F}(tf.func, tf.arity, tf.ir, tf.tape,
-               tf.counter, tf.bindings, :none)
+               tf.counter, tf.bindings, TypedVar{Any}(:none))
     end
 end
 
@@ -65,7 +69,7 @@ val(x) = x
 val(x::GlobalRef) = getproperty(x.mod, x.name)
 val(x::QuoteNode) = eval(x)
 val(x::TapedFunction) = x.func
-result(t::TapedFunction) = val(t.bindings[t.retval])
+result(t::TapedFunction) = _lookup(t, t.retval)
 
 function (tf::TapedFunction)(args...; callback=nothing)
     # set args
@@ -82,7 +86,7 @@ function (tf::TapedFunction)(args...; callback=nothing)
         ins = tf.tape[tf.counter]
         ins(tf)
         callback !== nothing && callback()
-        tf.retval !== :none && break
+        tf.retval.id !== :none && break
     end
     return result(tf)
 end
@@ -98,9 +102,9 @@ function Base.show(io::IO, tf::TapedFunction)
     println(buf, "------------------")
     println(buf, tf.ir)
     println(buf, "------------------")
-    println(buf, "* .tape =>")
+    println(buf, "* .bindings =>")
     println(buf, "------------------")
-    println(buf, tf.tape)
+    println(buf, tf.bindings)
     println(buf, "------------------")
     print(io, String(take!(buf)))
 end
@@ -118,8 +122,9 @@ end
 
 
 _lookup(tf::TapedFunction, v) = v
-_lookup(tf::TapedFunction, v::Symbol) = tf.bindings[v]
+_lookup(tf::TapedFunction, v::TypedVar{T}) where T = tf.bindings[v.id]::T
 _update_var!(tf::TapedFunction, v::Symbol, c) = tf.bindings[v] = c
+_update_var!(tf::TapedFunction, v::TypedVar{T}, c) where T = tf.bindings[v.id] = c
 
 function (instr::Instruction{F})(tf::TapedFunction) where F
     # catch run-time exceptions / errors.
@@ -138,8 +143,8 @@ function (instr::Instruction{F})(tf::TapedFunction) where F
 end
 
 function (instr::GotoInstruction)(tf::TapedFunction)
-    cond = instr.condition === :_true ? true :
-        instr.condition === :_false ? false :
+    cond = instr.condition.id === :_true ? true :
+        instr.condition.id === :_false ? false :
         val(_lookup(tf, instr.condition))
 
     if cond
@@ -158,6 +163,8 @@ end
 
 _accurate_typeof(v) = typeof(v)
 _accurate_typeof(::Type{V}) where V = Type{V}
+_loose_type(t) = t
+_loose_type(::Type{Type{T}}) where T = isa(T, DataType) ? Type{T} : typeof(T)
 
 """
     __new__(T, args...)
@@ -172,20 +179,24 @@ end
 
 ## Translation: CodeInfo -> Tape
 
-function bind_var!(var, bindings::Dict{Symbol, Any}) # for literal constants
+function bind_var!(var, bindings::Dict{Symbol, Any}, ir::Core.CodeInfo) # for literal constants
     id = gensym()
     bindings[id] = var
-    return id
+    return TypedVar{typeof(var)}(id)
 end
-bind_var!(var::Core.SSAValue, bindings::Dict{Symbol, Any}) =
-    bind_var!(Symbol(var.id), bindings)
-bind_var!(var::Core.TypedSlot, bindings::Dict{Symbol, Any}) =
-    bind_var!(Symbol(:_, var.id), bindings)
-bind_var!(var::Core.SlotNumber, bindings::Dict{Symbol, Any}) =
-    bind_var!(Symbol(:_, var.id), bindings)
-function bind_var!(var::Symbol, bindings::Dict{Symbol, Any})
+bind_var!(var::Core.SSAValue, bindings::Dict{Symbol, Any}, ir::Core.CodeInfo) =
+    bind_var!(Symbol(var.id), bindings, ir.ssavaluetypes[var.id])
+bind_var!(var::Core.TypedSlot, bindings::Dict{Symbol, Any}, ir::Core.CodeInfo) =
+    bind_var!(Symbol(:_, var.id), bindings, ir.slottypes[var.id])
+bind_var!(var::Core.SlotNumber, bindings::Dict{Symbol, Any}, ir::Core.CodeInfo) =
+    bind_var!(Symbol(:_, var.id), bindings, ir.slottypes[var.id])
+bind_var!(var::Symbol, boxes::Dict{Symbol, Any}, c::Core.Const) =
+    bind_var!(var, boxes, _loose_type(Type{c.val}))
+bind_var!(var::Symbol, boxes::Dict{Symbol, Any}, c::Core.PartialStruct) =
+    bind_var!(var, boxes, _loose_type(c.typ))
+function bind_var!(var::Symbol, bindings::Dict{Symbol, Any}, ::Type{T}) where T
     get!(bindings, var, nothing)
-    return var
+    TypedVar{T}(var)
 end
 
 function translate!(tape::Vector{FunctionWrapper{Nothing, Tuple{TapedFunction}}}, ir::Core.CodeInfo)
@@ -205,30 +216,30 @@ function translate!!(var::IRVar, line::Core.NewvarNode,
                      bindings::Dict{Symbol, Any}, @nospecialize(ir))
     # use a noop to ensure the 1-to-1 mapping from ir.code to instructions
     # on tape. see GotoInstruction.dest.
-    return GotoInstruction(:_true, 0)
+    return GotoInstruction(TypedVar{Bool}(:_true), 0)
 end
 
 function translate!!(var::IRVar, line::GlobalRef,
-                     bindings::Dict{Symbol, Any}, @nospecialize(ir))
-    return Instruction(() -> val(line), (), bind_var!(var, bindings))
+                     bindings::Dict{Symbol, Any}, ir)
+    return Instruction(() -> val(line), (), bind_var!(var, bindings, ir))
 end
 
 function translate!!(var::IRVar, line::Core.SlotNumber,
-                     bindings::Dict{Symbol, Any}, @nospecialize(ir))
-    return Instruction(identity, (bind_var!(line, bindings),), bind_var!(var, bindings))
+                     bindings::Dict{Symbol, Any}, ir)
+    return Instruction(identity, (bind_var!(line, bindings, ir),), bind_var!(var, bindings, ir))
 end
 
 function translate!!(var::IRVar, line::Core.TypedSlot,
-                     bindings::Dict{Symbol, Any}, @nospecialize(ir))
-    input_box = bind_var!(Core.SlotNumber(line.id), bindings)
-    return Instruction(identity, (input_box,), bind_var!(var, bindings))
+                     bindings::Dict{Symbol, Any}, ir)
+    input_box = bind_var!(Core.SlotNumber(line.id), bindings, ir)
+    return Instruction(identity, (input_box,), bind_var!(var, bindings, ir))
 end
 
 function translate!!(var::IRVar, line::Core.GotoIfNot,
-                     bindings::Dict{Symbol, Any}, @nospecialize(ir))
-    _cond = bind_var!(line.cond, bindings)
+                     bindings::Dict{Symbol, Any}, ir)
+    _cond = bind_var!(line.cond, bindings, ir)
     cond = if isa(_cond, Bool)
-        _cond ? :_true : :_false
+        TypedVar{Bool}(_cond ? :_true : :_false)
     else
         _cond
     end
@@ -237,21 +248,21 @@ end
 
 function translate!!(var::IRVar, line::Core.GotoNode,
                      bindings::Dict{Symbol, Any}, @nospecialize(ir))
-    return GotoInstruction(:_false, line.label)
+    return GotoInstruction(TypedVar{Bool}(:_false), line.label)
 end
 
 function translate!!(var::IRVar, line::Core.ReturnNode,
-                     bindings::Dict{Symbol, Any}, @nospecialize(ir))
-    return ReturnInstruction(bind_var!(line.val, bindings))
+                     bindings::Dict{Symbol, Any}, ir)
+    return ReturnInstruction(bind_var!(line.val, bindings, ir))
 end
 
 function translate!!(var::IRVar, line::Expr,
                      bindings::Dict{Symbol, Any}, ir::Core.CodeInfo)
     head = line.head
-    _bind_fn = (x) -> bind_var!(x, bindings)
+    _bind_fn = (x) -> bind_var!(x, bindings, ir)
     if head === :new
         args = map(_bind_fn, line.args)
-        return Instruction(__new__, args |> Tuple, bind_var!(var, bindings))
+        return Instruction(__new__, args |> Tuple, _bind_fn(var))
     elseif head === :call
         args = map(_bind_fn, line.args)
         # args[1] is the function
@@ -261,7 +272,7 @@ function translate!!(var::IRVar, line::Expr,
         else # isa(func, GlobalRef) or a var?
             func = args[1] # a var(box)
         end
-        return Instruction(func, args[2:end] |> Tuple, bind_var!(var, bindings))
+        return Instruction(func, args[2:end] |> Tuple, _bind_fn(var))
     elseif head === :(=)
         # line.args[1] (the left hand side) is a SlotNumber, and it should be the output
         lhs = line.args[1]

From e929b234fe176b9fb3ab14dba6bc3be4affcacec Mon Sep 17 00:00:00 2001
From: KDr2 <zhuo.dev@gmail.com>
Date: Sat, 5 Mar 2022 12:05:38 +0800
Subject: [PATCH 04/15] toggle FunctionWrapper usage

---
 src/tapedfunction.jl | 27 ++++++++++++++++++++-------
 1 file changed, 20 insertions(+), 7 deletions(-)

diff --git a/src/tapedfunction.jl b/src/tapedfunction.jl
index 69128720..e619c7c9 100644
--- a/src/tapedfunction.jl
+++ b/src/tapedfunction.jl
@@ -1,6 +1,7 @@
 ## Instruction and TapedFunction
 
 abstract type AbstractInstruction end
+const RawTape = Vector{AbstractInstruction}
 
 struct TypedVar{T}
     id::Symbol
@@ -32,7 +33,8 @@ mutable struct TapedFunction{F}
     func::F # maybe a function, a constructor, or a callable object
     arity::Int
     ir::Core.CodeInfo
-    tape::Vector{FunctionWrapper{Nothing, Tuple{TapedFunction}}}
+    tape::RawTape
+    unified_tape::Vector{FunctionWrapper{Nothing, Tuple{TapedFunction}}}
     counter::Int
     bindings::Dict{Symbol, Any}
     retval::TypedVar
@@ -49,20 +51,29 @@ mutable struct TapedFunction{F}
         end
 
         ir = CodeInfoTools.code_inferred(f, args_type...)
-        tape = Vector{FunctionWrapper{Nothing, Tuple{TapedFunction}}}()
+        tape = RawTape()
+        utape = Vector{FunctionWrapper{Nothing, Tuple{TapedFunction}}}()
         bindings = translate!(tape, ir)
 
-        tf = new{F}(f, length(args), ir, tape, 1, bindings, TypedVar{Any}(:none))
+        tf = new{F}(f, length(args), ir, tape, utape, 1, bindings, TypedVar{Any}(:none))
         TRCache[cache_key] = tf # set cache
+        # unify!(tf)
         return tf
     end
 
     function TapedFunction(tf::TapedFunction{F}) where {F}
-        new{F}(tf.func, tf.arity, tf.ir, tf.tape,
+        new{F}(tf.func, tf.arity, tf.ir, tf.tape, tf.unified_tape,
                tf.counter, tf.bindings, TypedVar{Any}(:none))
     end
 end
 
+function unify!(tf::TapedFunction)
+    length(tf.tape) == length(tf.unified_tape) && return
+    for ins in tf.tape
+        push!(tf.unified_tape, FunctionWrapper{Nothing, Tuple{TapedFunction}}(ins))
+    end
+end
+
 const TRCache = LRU{Tuple, TapedFunction}(maxsize=10)
 
 val(x) = x
@@ -82,8 +93,10 @@ function (tf::TapedFunction)(args...; callback=nothing)
     end
 
     # run the raw tape
+    tape = length(tf.tape) == length(tf.unified_tape) ?
+        tf.unified_tape : tf.tape
     while true
-        ins = tf.tape[tf.counter]
+        ins = tape[tf.counter]
         ins(tf)
         callback !== nothing && callback()
         tf.retval.id !== :none && break
@@ -199,13 +212,13 @@ function bind_var!(var::Symbol, bindings::Dict{Symbol, Any}, ::Type{T}) where T
     TypedVar{T}(var)
 end
 
-function translate!(tape::Vector{FunctionWrapper{Nothing, Tuple{TapedFunction}}}, ir::Core.CodeInfo)
+function translate!(tape::RawTape, ir::Core.CodeInfo)
     bindings = Dict{Symbol, Any}()
 
     for (idx, line) in enumerate(ir.code)
         isa(line, Core.Const) && (line = line.val) # unbox Core.Const
         ins = translate!!(Core.SSAValue(idx), line, bindings, ir)
-        push!(tape, FunctionWrapper{Nothing, Tuple{TapedFunction}}(ins))
+        push!(tape, ins)
     end
     return bindings
 end

From 876b51470f74fcf80303871b72823ab3ba0d4042 Mon Sep 17 00:00:00 2001
From: KDr2 <zhuo.dev@gmail.com>
Date: Sat, 5 Mar 2022 20:23:12 +0800
Subject: [PATCH 05/15] add return value type as TapedFunction's static
 parameter

---
 src/tapedfunction.jl | 27 +++++++++++++++++++++------
 test/tapedtask.jl    |  2 +-
 test/tf.jl           |  4 ++--
 3 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/src/tapedfunction.jl b/src/tapedfunction.jl
index e619c7c9..634a26f6 100644
--- a/src/tapedfunction.jl
+++ b/src/tapedfunction.jl
@@ -29,7 +29,7 @@ struct ReturnInstruction{T} <: AbstractInstruction
     arg::TypedVar{T}
 end
 
-mutable struct TapedFunction{F}
+mutable struct TapedFunction{F, T}
     func::F # maybe a function, a constructor, or a callable object
     arity::Int
     ir::Core.CodeInfo
@@ -37,7 +37,7 @@ mutable struct TapedFunction{F}
     unified_tape::Vector{FunctionWrapper{Nothing, Tuple{TapedFunction}}}
     counter::Int
     bindings::Dict{Symbol, Any}
-    retval::TypedVar
+    retval::TypedVar{T}
 
     function TapedFunction(f::F, args...; cache=false) where {F}
         args_type = _accurate_typeof.(args)
@@ -55,15 +55,17 @@ mutable struct TapedFunction{F}
         utape = Vector{FunctionWrapper{Nothing, Tuple{TapedFunction}}}()
         bindings = translate!(tape, ir)
 
-        tf = new{F}(f, length(args), ir, tape, utape, 1, bindings, TypedVar{Any}(:none))
+        T = isa(ir.rettype, Core.Const) ? typeof(ir.rettype.val) : ir.rettype
+        tf = new{F, T}(f, length(args), ir, tape, utape, 1,
+                                bindings, TypedVar{T}(:none))
         TRCache[cache_key] = tf # set cache
         # unify!(tf)
         return tf
     end
 
-    function TapedFunction(tf::TapedFunction{F}) where {F}
-        new{F}(tf.func, tf.arity, tf.ir, tf.tape, tf.unified_tape,
-               tf.counter, tf.bindings, TypedVar{Any}(:none))
+    function TapedFunction(tf::TapedFunction{F, T}) where {F, T}
+        new{F, T}(tf.func, tf.arity, tf.ir, tf.tape, tf.unified_tape,
+                  tf.counter, tf.bindings, tf.retval)
     end
 end
 
@@ -122,6 +124,19 @@ function Base.show(io::IO, tf::TapedFunction)
     print(io, String(take!(buf)))
 end
 
+function Base.show(io::IO, rtape::RawTape)
+    buf = IOBuffer()
+    print(buf, length(rtape), "-element RawTape")
+    isempty(rtape) || println(buf, ":")
+    i = 1
+    for instr in rtape
+        print(buf, "\t", i, " => ")
+        show(buf, instr)
+        i += 1
+    end
+    print(io, String(take!(buf)))
+end
+
 ## methods for Instruction
 Base.show(io::IO, instr::AbstractInstruction) = println(io, "A ", typeof(instr))
 
diff --git a/test/tapedtask.jl b/test/tapedtask.jl
index 041431ab..9a45824c 100644
--- a/test/tapedtask.jl
+++ b/test/tapedtask.jl
@@ -18,7 +18,7 @@
         @test consume(ttask) == 2
         @test consume(ttask) == 3
 
-        @inferred Libtask.TapedFunction(f)
+        # @inferred Libtask.TapedFunction(f)
     end
 
     # Test case 2: heap allocated objects are shallowly copied.
diff --git a/test/tf.jl b/test/tf.jl
index 3192460b..b4af5dab 100644
--- a/test/tf.jl
+++ b/test/tf.jl
@@ -11,7 +11,7 @@ using Libtask
         tf = Libtask.TapedFunction(S, 1, 2)
         s1 = tf(1, 2)
         @test s1.i == 3
-        # newins = findall(x -> isa(x, Libtask.Instruction{typeof(Libtask.__new__)}), tf.tape)
-        # @test length(newins) == 1
+        newins = findall(x -> isa(x, Libtask.Instruction{typeof(Libtask.__new__)}), tf.tape)
+        @test length(newins) == 1
     end
 end

From f904035e56cae2f3e12ade58baa6c03cf85ea58f Mon Sep 17 00:00:00 2001
From: KDr2 <zhuo.dev@gmail.com>
Date: Sun, 6 Mar 2022 08:06:27 +0800
Subject: [PATCH 06/15] don't use inferred return type

---
 src/tapedfunction.jl | 15 +++++++--------
 test/tapedtask.jl    |  2 +-
 2 files changed, 8 insertions(+), 9 deletions(-)

diff --git a/src/tapedfunction.jl b/src/tapedfunction.jl
index 634a26f6..02437608 100644
--- a/src/tapedfunction.jl
+++ b/src/tapedfunction.jl
@@ -29,7 +29,7 @@ struct ReturnInstruction{T} <: AbstractInstruction
     arg::TypedVar{T}
 end
 
-mutable struct TapedFunction{F, T}
+mutable struct TapedFunction{F}
     func::F # maybe a function, a constructor, or a callable object
     arity::Int
     ir::Core.CodeInfo
@@ -37,7 +37,7 @@ mutable struct TapedFunction{F, T}
     unified_tape::Vector{FunctionWrapper{Nothing, Tuple{TapedFunction}}}
     counter::Int
     bindings::Dict{Symbol, Any}
-    retval::TypedVar{T}
+    retval::TypedVar{<:Any}
 
     function TapedFunction(f::F, args...; cache=false) where {F}
         args_type = _accurate_typeof.(args)
@@ -55,17 +55,16 @@ mutable struct TapedFunction{F, T}
         utape = Vector{FunctionWrapper{Nothing, Tuple{TapedFunction}}}()
         bindings = translate!(tape, ir)
 
-        T = isa(ir.rettype, Core.Const) ? typeof(ir.rettype.val) : ir.rettype
-        tf = new{F, T}(f, length(args), ir, tape, utape, 1,
-                                bindings, TypedVar{T}(:none))
+        tf = new{F}(f, length(args), ir, tape, utape, 1,
+                                bindings, TypedVar{Any}(:none))
         TRCache[cache_key] = tf # set cache
         # unify!(tf)
         return tf
     end
 
-    function TapedFunction(tf::TapedFunction{F, T}) where {F, T}
-        new{F, T}(tf.func, tf.arity, tf.ir, tf.tape, tf.unified_tape,
-                  tf.counter, tf.bindings, tf.retval)
+    function TapedFunction(tf::TapedFunction{F}) where {F}
+        new{F}(tf.func, tf.arity, tf.ir, tf.tape, tf.unified_tape,
+                  tf.counter, tf.bindings, TypedVar{Any}(:none))
     end
 end
 
diff --git a/test/tapedtask.jl b/test/tapedtask.jl
index 9a45824c..041431ab 100644
--- a/test/tapedtask.jl
+++ b/test/tapedtask.jl
@@ -18,7 +18,7 @@
         @test consume(ttask) == 2
         @test consume(ttask) == 3
 
-        # @inferred Libtask.TapedFunction(f)
+        @inferred Libtask.TapedFunction(f)
     end
 
     # Test case 2: heap allocated objects are shallowly copied.

From 316d9f8fcdad377ec14c16dc07a7ffdedd976970 Mon Sep 17 00:00:00 2001
From: KDr2 <zhuo.dev@gmail.com>
Date: Tue, 8 Mar 2022 18:40:53 +0800
Subject: [PATCH 07/15] fix ir

---
 src/tapedfunction.jl | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/src/tapedfunction.jl b/src/tapedfunction.jl
index 02437608..9ca5b2da 100644
--- a/src/tapedfunction.jl
+++ b/src/tapedfunction.jl
@@ -51,6 +51,7 @@ mutable struct TapedFunction{F}
         end
 
         ir = CodeInfoTools.code_inferred(f, args_type...)
+        fix_ir(ir)
         tape = RawTape()
         utape = Vector{FunctionWrapper{Nothing, Tuple{TapedFunction}}}()
         bindings = translate!(tape, ir)
@@ -187,6 +188,13 @@ end
 
 
 ## internal functions
+function fix_ir(ir::Core.CodeInfo)
+    if Base.JLOptions().code_coverage != 0 &&
+        length(ir.ssavaluetypes) > length(ir.code) &&
+        ir.ssavaluetypes[1] === Nothing
+        popfirst!(ir.ssavaluetypes)
+    end
+end
 
 _accurate_typeof(v) = typeof(v)
 _accurate_typeof(::Type{V}) where V = Type{V}

From 09940472a7e941cb5aa48e19561023d1afb1fbeb Mon Sep 17 00:00:00 2001
From: KDr2 <zhuo.dev@gmail.com>
Date: Tue, 8 Mar 2022 18:57:56 +0800
Subject: [PATCH 08/15] disable coverage when run tests

---
 .github/workflows/Testing.yaml | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/.github/workflows/Testing.yaml b/.github/workflows/Testing.yaml
index 48d03b62..a54a7c83 100644
--- a/.github/workflows/Testing.yaml
+++ b/.github/workflows/Testing.yaml
@@ -42,3 +42,5 @@ jobs:
             ${{ runner.os }}-
       - uses: julia-actions/julia-buildpkg@latest
       - uses: julia-actions/julia-runtest@latest
+        with:
+          coverage: false

From 0fc4bdf22755b36eb84e3adc1175b2edf41557db Mon Sep 17 00:00:00 2001
From: KDr2 <zhuo.dev@gmail.com>
Date: Tue, 8 Mar 2022 19:43:52 +0800
Subject: [PATCH 09/15] Revert "fix ir"

This reverts commit 316d9f8fcdad377ec14c16dc07a7ffdedd976970.
---
 src/tapedfunction.jl | 8 --------
 1 file changed, 8 deletions(-)

diff --git a/src/tapedfunction.jl b/src/tapedfunction.jl
index 9ca5b2da..02437608 100644
--- a/src/tapedfunction.jl
+++ b/src/tapedfunction.jl
@@ -51,7 +51,6 @@ mutable struct TapedFunction{F}
         end
 
         ir = CodeInfoTools.code_inferred(f, args_type...)
-        fix_ir(ir)
         tape = RawTape()
         utape = Vector{FunctionWrapper{Nothing, Tuple{TapedFunction}}}()
         bindings = translate!(tape, ir)
@@ -188,13 +187,6 @@ end
 
 
 ## internal functions
-function fix_ir(ir::Core.CodeInfo)
-    if Base.JLOptions().code_coverage != 0 &&
-        length(ir.ssavaluetypes) > length(ir.code) &&
-        ir.ssavaluetypes[1] === Nothing
-        popfirst!(ir.ssavaluetypes)
-    end
-end
 
 _accurate_typeof(v) = typeof(v)
 _accurate_typeof(::Type{V}) where V = Type{V}

From 44b0bfaca8be6aa504243985fbe854177792cbf7 Mon Sep 17 00:00:00 2001
From: KDr2 <zhuo.dev@gmail.com>
Date: Wed, 16 Mar 2022 09:51:55 +0800
Subject: [PATCH 10/15] put getter/setter into box

---
 src/tapedfunction.jl | 100 ++++++++++++++++++++++++-------------------
 1 file changed, 55 insertions(+), 45 deletions(-)

diff --git a/src/tapedfunction.jl b/src/tapedfunction.jl
index 02437608..26f41d75 100644
--- a/src/tapedfunction.jl
+++ b/src/tapedfunction.jl
@@ -3,32 +3,6 @@
 abstract type AbstractInstruction end
 const RawTape = Vector{AbstractInstruction}
 
-struct TypedVar{T}
-    id::Symbol
-end
-
-"""
-    Instruction
-
-An `Instruction` stands for a function call
-"""
-struct Instruction{F, N, TO} <: AbstractInstruction
-    func::F
-    input::NTuple{N, TypedVar{<:Any}}
-    output::TypedVar{TO}
-end
-
-struct GotoInstruction{T} <: AbstractInstruction
-    condition::TypedVar{T}
-    # we enusre a 1-to-1 mapping between ir.code and instruction
-    # so here we can use the index directly.
-    dest::Int
-end
-
-struct ReturnInstruction{T} <: AbstractInstruction
-    arg::TypedVar{T}
-end
-
 mutable struct TapedFunction{F}
     func::F # maybe a function, a constructor, or a callable object
     arity::Int
@@ -37,7 +11,7 @@ mutable struct TapedFunction{F}
     unified_tape::Vector{FunctionWrapper{Nothing, Tuple{TapedFunction}}}
     counter::Int
     bindings::Dict{Symbol, Any}
-    retval::TypedVar{<:Any}
+    retval::Symbol
 
     function TapedFunction(f::F, args...; cache=false) where {F}
         args_type = _accurate_typeof.(args)
@@ -55,8 +29,7 @@ mutable struct TapedFunction{F}
         utape = Vector{FunctionWrapper{Nothing, Tuple{TapedFunction}}}()
         bindings = translate!(tape, ir)
 
-        tf = new{F}(f, length(args), ir, tape, utape, 1,
-                                bindings, TypedVar{Any}(:none))
+        tf = new{F}(f, length(args), ir, tape, utape, 1, bindings, :none)
         TRCache[cache_key] = tf # set cache
         # unify!(tf)
         return tf
@@ -64,10 +37,12 @@ mutable struct TapedFunction{F}
 
     function TapedFunction(tf::TapedFunction{F}) where {F}
         new{F}(tf.func, tf.arity, tf.ir, tf.tape, tf.unified_tape,
-                  tf.counter, tf.bindings, TypedVar{Any}(:none))
+               tf.counter, tf.bindings, :none)
     end
 end
 
+const TRCache = LRU{Tuple, TapedFunction}(maxsize=10)
+
 function unify!(tf::TapedFunction)
     length(tf.tape) == length(tf.unified_tape) && return
     for ins in tf.tape
@@ -75,13 +50,54 @@ function unify!(tf::TapedFunction)
     end
 end
 
-const TRCache = LRU{Tuple, TapedFunction}(maxsize=10)
+struct Box{T}
+    id::Symbol
+    get # ::FunctionWrapper{T, Tuple{TapedFunction, Symbol}}
+    set # ::FunctionWrapper{Nothing, Tuple{TapedFunction, Symbol, T}}
+
+    function Box{T}(id::Symbol) where T
+        return new(id, _inner_getter, _inner_setter)
+        # return new(id,
+        # FunctionWrapper{T, Tuple{TapedFunction, Symbol}}(_inner_getter),
+        # FunctionWrapper{Nothing, Tuple{TapedFunction, Symbol, T}}(_inner_setter))
+    end
+end
+
+_inner_getter(tf::TapedFunction, v::Symbol) = tf.bindings[v]
+_inner_setter(tf::TapedFunction, v::Symbol, c) = tf.bindings[v] = c
+_lookup(tf::TapedFunction, v) = v
+_lookup(tf::TapedFunction, v::Box{T}) where T = v.get(tf, v.id)
+_update_var!(tf::TapedFunction, v::Symbol, c) = tf.bindings[v] = c
+_update_var!(tf::TapedFunction, v::Box{T}, c::T) where T = v.set(tf, v.id, c)
+
+"""
+    Instruction
+
+An `Instruction` stands for a function call
+"""
+struct Instruction{F, N, TO} <: AbstractInstruction
+    func::F
+    input::NTuple{N, Box{<:Any}}
+    output::Box{TO}
+end
+
+struct GotoInstruction{T} <: AbstractInstruction
+    condition::Box{T}
+    # we enusre a 1-to-1 mapping between ir.code and instruction
+    # so here we can use the index directly.
+    dest::Int
+end
+
+struct ReturnInstruction{T} <: AbstractInstruction
+    arg::Box{T}
+end
+
 
 val(x) = x
 val(x::GlobalRef) = getproperty(x.mod, x.name)
 val(x::QuoteNode) = eval(x)
 val(x::TapedFunction) = x.func
-result(t::TapedFunction) = _lookup(t, t.retval)
+result(t::TapedFunction) = t.bindings[t.retval]
 
 function (tf::TapedFunction)(args...; callback=nothing)
     # set args
@@ -100,7 +116,7 @@ function (tf::TapedFunction)(args...; callback=nothing)
         ins = tape[tf.counter]
         ins(tf)
         callback !== nothing && callback()
-        tf.retval.id !== :none && break
+        tf.retval !== :none && break
     end
     return result(tf)
 end
@@ -147,12 +163,6 @@ function Base.show(io::IO, instr::GotoInstruction)
     println(io, "GotoInstruction(", instr.condition, ", dest=", instr.dest, ")")
 end
 
-
-_lookup(tf::TapedFunction, v) = v
-_lookup(tf::TapedFunction, v::TypedVar{T}) where T = tf.bindings[v.id]::T
-_update_var!(tf::TapedFunction, v::Symbol, c) = tf.bindings[v] = c
-_update_var!(tf::TapedFunction, v::TypedVar{T}, c) where T = tf.bindings[v.id] = c
-
 function (instr::Instruction{F})(tf::TapedFunction) where F
     # catch run-time exceptions / errors.
     try
@@ -182,7 +192,7 @@ function (instr::GotoInstruction)(tf::TapedFunction)
 end
 
 function (instr::ReturnInstruction)(tf::TapedFunction)
-    tf.retval = instr.arg
+    tf.retval = instr.arg.id
 end
 
 
@@ -209,7 +219,7 @@ end
 function bind_var!(var, bindings::Dict{Symbol, Any}, ir::Core.CodeInfo) # for literal constants
     id = gensym()
     bindings[id] = var
-    return TypedVar{typeof(var)}(id)
+    Box{typeof(var)}(id)
 end
 bind_var!(var::Core.SSAValue, bindings::Dict{Symbol, Any}, ir::Core.CodeInfo) =
     bind_var!(Symbol(var.id), bindings, ir.ssavaluetypes[var.id])
@@ -223,7 +233,7 @@ bind_var!(var::Symbol, boxes::Dict{Symbol, Any}, c::Core.PartialStruct) =
     bind_var!(var, boxes, _loose_type(c.typ))
 function bind_var!(var::Symbol, bindings::Dict{Symbol, Any}, ::Type{T}) where T
     get!(bindings, var, nothing)
-    TypedVar{T}(var)
+    Box{T}(var)
 end
 
 function translate!(tape::RawTape, ir::Core.CodeInfo)
@@ -243,7 +253,7 @@ function translate!!(var::IRVar, line::Core.NewvarNode,
                      bindings::Dict{Symbol, Any}, @nospecialize(ir))
     # use a noop to ensure the 1-to-1 mapping from ir.code to instructions
     # on tape. see GotoInstruction.dest.
-    return GotoInstruction(TypedVar{Bool}(:_true), 0)
+    return GotoInstruction(Box{Bool}(:_true), 0)
 end
 
 function translate!!(var::IRVar, line::GlobalRef,
@@ -266,7 +276,7 @@ function translate!!(var::IRVar, line::Core.GotoIfNot,
                      bindings::Dict{Symbol, Any}, ir)
     _cond = bind_var!(line.cond, bindings, ir)
     cond = if isa(_cond, Bool)
-        TypedVar{Bool}(_cond ? :_true : :_false)
+        Box{Bool}(_cond ? :_true : :_false)
     else
         _cond
     end
@@ -275,7 +285,7 @@ end
 
 function translate!!(var::IRVar, line::Core.GotoNode,
                      bindings::Dict{Symbol, Any}, @nospecialize(ir))
-    return GotoInstruction(TypedVar{Bool}(:_false), line.label)
+    return GotoInstruction(Box{Bool}(:_false), line.label)
 end
 
 function translate!!(var::IRVar, line::Core.ReturnNode,

From 858b01c03c32d78cdf9fd24fcfb0271c3bf57f02 Mon Sep 17 00:00:00 2001
From: KDr2 <zhuo.dev@gmail.com>
Date: Thu, 17 Mar 2022 11:16:17 +0800
Subject: [PATCH 11/15] add TypedFunction

---
 src/tapedfunction.jl | 23 +++++++++++++++++------
 1 file changed, 17 insertions(+), 6 deletions(-)

diff --git a/src/tapedfunction.jl b/src/tapedfunction.jl
index 26f41d75..b0b5b7e1 100644
--- a/src/tapedfunction.jl
+++ b/src/tapedfunction.jl
@@ -50,16 +50,27 @@ function unify!(tf::TapedFunction)
     end
 end
 
+# const TypedFunction = FunctionWrapper
+mutable struct TypedFunction{OT, IT<:Tuple}
+    func::Function
+    retval::OT
+    TypedFunction{OT, IT}(f::Function) where {OT, IT<:Tuple} = new{OT, IT}(f)
+end
+function (f::TypedFunction{OT, IT})(args...) where {OT, IT<:Tuple}
+    output = f.func(args...)
+    OT === Nothing ? (f.retval = nothing) : (f.retval = output)
+    return f.retval
+end
+
 struct Box{T}
     id::Symbol
-    get # ::FunctionWrapper{T, Tuple{TapedFunction, Symbol}}
-    set # ::FunctionWrapper{Nothing, Tuple{TapedFunction, Symbol, T}}
+    get::TypedFunction{T, Tuple{TapedFunction, Symbol}}
+    set::TypedFunction{Nothing, Tuple{TapedFunction, Symbol, T}}
 
     function Box{T}(id::Symbol) where T
-        return new(id, _inner_getter, _inner_setter)
-        # return new(id,
-        # FunctionWrapper{T, Tuple{TapedFunction, Symbol}}(_inner_getter),
-        # FunctionWrapper{Nothing, Tuple{TapedFunction, Symbol, T}}(_inner_setter))
+        return new(id,
+                   TypedFunction{T, Tuple{TapedFunction, Symbol}}(_inner_getter),
+                   TypedFunction{Nothing, Tuple{TapedFunction, Symbol, T}}(_inner_setter))
     end
 end
 

From beb6b5280f6cb4a5aa2498773bb8265106b3e3e6 Mon Sep 17 00:00:00 2001
From: KDr2 <zhuo.dev@gmail.com>
Date: Thu, 17 Mar 2022 17:28:36 +0800
Subject: [PATCH 12/15] inline functions

---
 src/tapedfunction.jl | 22 +++++++++++-----------
 1 file changed, 11 insertions(+), 11 deletions(-)

diff --git a/src/tapedfunction.jl b/src/tapedfunction.jl
index b0b5b7e1..7134745b 100644
--- a/src/tapedfunction.jl
+++ b/src/tapedfunction.jl
@@ -51,15 +51,15 @@ function unify!(tf::TapedFunction)
 end
 
 # const TypedFunction = FunctionWrapper
-mutable struct TypedFunction{OT, IT<:Tuple}
+struct TypedFunction{OT, IT<:Tuple}
     func::Function
-    retval::OT
-    TypedFunction{OT, IT}(f::Function) where {OT, IT<:Tuple} = new{OT, IT}(f)
+    retval::Ref{OT}
+    TypedFunction{OT, IT}(f::Function) where {OT, IT<:Tuple} = new{OT, IT}(f, Ref{OT}())
 end
 function (f::TypedFunction{OT, IT})(args...) where {OT, IT<:Tuple}
     output = f.func(args...)
-    OT === Nothing ? (f.retval = nothing) : (f.retval = output)
-    return f.retval
+    OT === Nothing ? (f.retval[] = nothing) : (f.retval[] = output)
+    return f.retval[]
 end
 
 struct Box{T}
@@ -74,12 +74,12 @@ struct Box{T}
     end
 end
 
-_inner_getter(tf::TapedFunction, v::Symbol) = tf.bindings[v]
-_inner_setter(tf::TapedFunction, v::Symbol, c) = tf.bindings[v] = c
-_lookup(tf::TapedFunction, v) = v
-_lookup(tf::TapedFunction, v::Box{T}) where T = v.get(tf, v.id)
-_update_var!(tf::TapedFunction, v::Symbol, c) = tf.bindings[v] = c
-_update_var!(tf::TapedFunction, v::Box{T}, c::T) where T = v.set(tf, v.id, c)
+@inline _inner_getter(tf::TapedFunction, v::Symbol) = tf.bindings[v]
+@inline _inner_setter(tf::TapedFunction, v::Symbol, c) = tf.bindings[v] = c
+@inline _lookup(tf::TapedFunction, v) = v
+@inline _lookup(tf::TapedFunction, v::Box{T}) where T = v.get(tf, v.id)
+@inline _update_var!(tf::TapedFunction, v::Symbol, c) = tf.bindings[v] = c
+@inline _update_var!(tf::TapedFunction, v::Box{T}, c::T) where T = v.set(tf, v.id, c)
 
 """
     Instruction

From f4229b958475d3d3c6d7ed717f1fde5a0f310bf0 Mon Sep 17 00:00:00 2001
From: KDr2 <zhuo.dev@gmail.com>
Date: Fri, 18 Mar 2022 06:32:08 +0800
Subject: [PATCH 13/15] compile tf

---
 src/tapedfunction.jl | 38 ++++++++++++++++++++------------------
 test/tapedtask.jl    |  2 +-
 test/tf.jl           | 17 +++++++++++++++++
 3 files changed, 38 insertions(+), 19 deletions(-)

diff --git a/src/tapedfunction.jl b/src/tapedfunction.jl
index 7134745b..7817d4ca 100644
--- a/src/tapedfunction.jl
+++ b/src/tapedfunction.jl
@@ -3,17 +3,17 @@
 abstract type AbstractInstruction end
 const RawTape = Vector{AbstractInstruction}
 
-mutable struct TapedFunction{F}
+mutable struct TapedFunction{F, TapeType}
     func::F # maybe a function, a constructor, or a callable object
     arity::Int
     ir::Core.CodeInfo
-    tape::RawTape
-    unified_tape::Vector{FunctionWrapper{Nothing, Tuple{TapedFunction}}}
+    tape::TapeType
     counter::Int
     bindings::Dict{Symbol, Any}
     retval::Symbol
 
-    function TapedFunction(f::F, args...; cache=false) where {F}
+    function TapedFunction(f, args...; cache=false)
+        F = typeof(f)
         args_type = _accurate_typeof.(args)
         cache_key = (f, args_type...)
 
@@ -26,39 +26,43 @@ mutable struct TapedFunction{F}
 
         ir = CodeInfoTools.code_inferred(f, args_type...)
         tape = RawTape()
-        utape = Vector{FunctionWrapper{Nothing, Tuple{TapedFunction}}}()
         bindings = translate!(tape, ir)
 
-        tf = new{F}(f, length(args), ir, tape, utape, 1, bindings, :none)
+        tf = new{F, RawTape}(f, length(args), ir, tape, 1, bindings, :none)
         TRCache[cache_key] = tf # set cache
-        # unify!(tf)
         return tf
     end
 
-    function TapedFunction(tf::TapedFunction{F}) where {F}
-        new{F}(tf.func, tf.arity, tf.ir, tf.tape, tf.unified_tape,
+    function TapedFunction{F, T0}(tf::TapedFunction{F, T1}) where {F, T0, T1}
+        new{F, T0}(tf.func, tf.arity, tf.ir, tf.tape,
                tf.counter, tf.bindings, :none)
     end
+
+    TapedFunction(tf::TapedFunction{F, T}) where {F, T} = TapedFunction{F, T}(tf)
 end
 
 const TRCache = LRU{Tuple, TapedFunction}(maxsize=10)
+const CompiledTape = Vector{FunctionWrapper{Nothing, Tuple{TapedFunction}}}
 
-function unify!(tf::TapedFunction)
-    length(tf.tape) == length(tf.unified_tape) && return
-    for ins in tf.tape
-        push!(tf.unified_tape, FunctionWrapper{Nothing, Tuple{TapedFunction}}(ins))
+function Base.convert(::Type{CompiledTape}, tape::RawTape)
+    ctape = CompiledTape(undef, length(tape))
+    for idx in 1:length(tape)
+        ctape[idx] = FunctionWrapper{Nothing, Tuple{TapedFunction}}(tape[idx])
     end
+    return ctape
 end
 
+compile(tf::TapedFunction{F, RawTape}) where {F} = TapedFunction{F, CompiledTape}(tf)
+
 # const TypedFunction = FunctionWrapper
 struct TypedFunction{OT, IT<:Tuple}
     func::Function
-    retval::Ref{OT}
+    retval::Base.RefValue{OT}
     TypedFunction{OT, IT}(f::Function) where {OT, IT<:Tuple} = new{OT, IT}(f, Ref{OT}())
 end
 function (f::TypedFunction{OT, IT})(args...) where {OT, IT<:Tuple}
     output = f.func(args...)
-    OT === Nothing ? (f.retval[] = nothing) : (f.retval[] = output)
+    f.retval[] = OT === Nothing ? nothing : output
     return f.retval[]
 end
 
@@ -121,10 +125,8 @@ function (tf::TapedFunction)(args...; callback=nothing)
     end
 
     # run the raw tape
-    tape = length(tf.tape) == length(tf.unified_tape) ?
-        tf.unified_tape : tf.tape
     while true
-        ins = tape[tf.counter]
+        ins = tf.tape[tf.counter]
         ins(tf)
         callback !== nothing && callback()
         tf.retval !== :none && break
diff --git a/test/tapedtask.jl b/test/tapedtask.jl
index 041431ab..9a45824c 100644
--- a/test/tapedtask.jl
+++ b/test/tapedtask.jl
@@ -18,7 +18,7 @@
         @test consume(ttask) == 2
         @test consume(ttask) == 3
 
-        @inferred Libtask.TapedFunction(f)
+        # @inferred Libtask.TapedFunction(f)
     end
 
     # Test case 2: heap allocated objects are shallowly copied.
diff --git a/test/tf.jl b/test/tf.jl
index b4af5dab..f1fd87c5 100644
--- a/test/tf.jl
+++ b/test/tf.jl
@@ -14,4 +14,21 @@ using Libtask
         newins = findall(x -> isa(x, Libtask.Instruction{typeof(Libtask.__new__)}), tf.tape)
         @test length(newins) == 1
     end
+
+    @testset "Compiled Tape" begin
+        function g(x, y)
+            if x>y
+                r= string(sin(x))
+            else
+                r= sin(x) * cos(y)
+            end
+            return r
+        end
+
+        tf = Libtask.TapedFunction(g, 1., 2.)
+        ctf = Libtask.compile(tf)
+        r = ctf(1., 2.)
+
+        @test typeof(r) === Float64
+    end
 end

From 454fb68ad3b4ee65878be6970fa0c4842bf4a889 Mon Sep 17 00:00:00 2001
From: KDr2 <zhuo.dev@gmail.com>
Date: Fri, 18 Mar 2022 06:42:06 +0800
Subject: [PATCH 14/15] fix parametiric type of tf

---
 src/tapedfunction.jl | 12 +++++++-----
 test/tapedtask.jl    |  2 +-
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/src/tapedfunction.jl b/src/tapedfunction.jl
index 7817d4ca..c1ac27bd 100644
--- a/src/tapedfunction.jl
+++ b/src/tapedfunction.jl
@@ -12,27 +12,29 @@ mutable struct TapedFunction{F, TapeType}
     bindings::Dict{Symbol, Any}
     retval::Symbol
 
-    function TapedFunction(f, args...; cache=false)
-        F = typeof(f)
+    function TapedFunction{F, T}(f::F, args...; cache=false) where {F, T}
         args_type = _accurate_typeof.(args)
         cache_key = (f, args_type...)
 
         if cache && haskey(TRCache, cache_key) # use cache
-            cached_tf = TRCache[cache_key]::TapedFunction{F}
+            cached_tf = TRCache[cache_key]::TapedFunction{F, T}
             tf = copy(cached_tf)
             tf.counter = 1
             return tf
         end
 
         ir = CodeInfoTools.code_inferred(f, args_type...)
-        tape = RawTape()
+        tape = T()
         bindings = translate!(tape, ir)
 
-        tf = new{F, RawTape}(f, length(args), ir, tape, 1, bindings, :none)
+        tf = new{F, T}(f, length(args), ir, tape, 1, bindings, :none)
         TRCache[cache_key] = tf # set cache
         return tf
     end
 
+    TapedFunction(f, args...; cache=false) =
+        TapedFunction{typeof(f), RawTape}(f, args...; cache=cache)
+
     function TapedFunction{F, T0}(tf::TapedFunction{F, T1}) where {F, T0, T1}
         new{F, T0}(tf.func, tf.arity, tf.ir, tf.tape,
                tf.counter, tf.bindings, :none)
diff --git a/test/tapedtask.jl b/test/tapedtask.jl
index 9a45824c..041431ab 100644
--- a/test/tapedtask.jl
+++ b/test/tapedtask.jl
@@ -18,7 +18,7 @@
         @test consume(ttask) == 2
         @test consume(ttask) == 3
 
-        # @inferred Libtask.TapedFunction(f)
+        @inferred Libtask.TapedFunction(f)
     end
 
     # Test case 2: heap allocated objects are shallowly copied.

From f79764d785c715e1b14512844766f840233f72b6 Mon Sep 17 00:00:00 2001
From: KDr2 <zhuo.dev@gmail.com>
Date: Fri, 18 Mar 2022 16:21:21 +0800
Subject: [PATCH 15/15] revise code

---
 src/tapedfunction.jl | 20 ++++++++------------
 1 file changed, 8 insertions(+), 12 deletions(-)

diff --git a/src/tapedfunction.jl b/src/tapedfunction.jl
index c1ac27bd..15f0143c 100644
--- a/src/tapedfunction.jl
+++ b/src/tapedfunction.jl
@@ -24,8 +24,7 @@ mutable struct TapedFunction{F, TapeType}
         end
 
         ir = CodeInfoTools.code_inferred(f, args_type...)
-        tape = T()
-        bindings = translate!(tape, ir)
+        bindings, tape = translate!(RawTape(), ir)
 
         tf = new{F, T}(f, length(args), ir, tape, 1, bindings, :none)
         TRCache[cache_key] = tf # set cache
@@ -62,6 +61,7 @@ struct TypedFunction{OT, IT<:Tuple}
     retval::Base.RefValue{OT}
     TypedFunction{OT, IT}(f::Function) where {OT, IT<:Tuple} = new{OT, IT}(f, Ref{OT}())
 end
+
 function (f::TypedFunction{OT, IT})(args...) where {OT, IT<:Tuple}
     output = f.func(args...)
     f.retval[] = OT === Nothing ? nothing : output
@@ -110,11 +110,11 @@ struct ReturnInstruction{T} <: AbstractInstruction
 end
 
 
-val(x) = x
-val(x::GlobalRef) = getproperty(x.mod, x.name)
-val(x::QuoteNode) = eval(x)
-val(x::TapedFunction) = x.func
-result(t::TapedFunction) = t.bindings[t.retval]
+@inline val(x) = x
+@inline val(x::GlobalRef) = getproperty(x.mod, x.name)
+@inline val(x::QuoteNode) = eval(x)
+@inline val(x::TapedFunction) = x.func
+@inline result(t::TapedFunction) = t.bindings[t.retval]
 
 function (tf::TapedFunction)(args...; callback=nothing)
     # set args
@@ -147,10 +147,6 @@ function Base.show(io::IO, tf::TapedFunction)
     println(buf, "------------------")
     println(buf, tf.ir)
     println(buf, "------------------")
-    println(buf, "* .bindings =>")
-    println(buf, "------------------")
-    println(buf, tf.bindings)
-    println(buf, "------------------")
     print(io, String(take!(buf)))
 end
 
@@ -259,7 +255,7 @@ function translate!(tape::RawTape, ir::Core.CodeInfo)
         ins = translate!!(Core.SSAValue(idx), line, bindings, ir)
         push!(tape, ins)
     end
-    return bindings
+    return (bindings, tape)
 end
 
 const IRVar = Union{Core.SSAValue, Core.SlotNumber}