<|start_filename|>src/fuzzer.jl<|end_filename|>

nbranch(t::Callsite) = nbranch(flatten(t))
function nbranch(stream)
    count = 0
    for (f, _, _) in stream
        if f == assert
            count+=1
        end
    end
    return count
end

function cut_at_nth_branch(stream, nth_branch)
    cut_stream = Tuple{Any, Any, Tuple}[]
    count = 0
    for (f, ret, args) in stream
        push!(cut_stream, (f, ret, args))
        if f == assert
            count += 1
            if count == nth_branch
                return cut_stream
            end
        end
    end
    return cut_stream
end

function force_and_cut(stream, nth_branch)
    cut_stream = cut_at_nth_branch(stream, nth_branch)
    f, ret, args = pop!(cut_stream)
    @assert f == assert
    @assert anything(ret) isa Bool
    push!(cut_stream, (f, Some(!anything(ret)), args))
end

function checkStream(stream)
    z3model = symbolic(stream)
    sat, model = runZ3(z3model)
    if sat
       inputs = parseZ3(model)
    else
       inputs = ()
    end

    return sat, inputs
end


import InteractiveUtils: subtypes

function concretize(T::Union{DataType, Union})
    ctypes = Any[]
    types = Any[T]

    while !isempty(types)
        T = pop!(types)
        if isconcretetype(T)
            push!(ctypes, T)
        elseif T isa Union
            push!(types, T.a)
            push!(types, T.b)
        else
            append!(types, subtypes(T))
        end
    end
    return ctypes
end

supported(T::DataType)= false
generate(T::DataType) = error("Can't generate values of $T")

supported(::Type{T}) where T<:INTEGERS = true
# generate(::Type{BigInt}) = rand(1:big(typemax(Int))^2)
generate(::Type{T}) where T<:INTEGERS = rand(T)

# UGLY HACK
function enumerateSupportedTypes(i)
    if i == 0
      return Int64
    elseif i == 1
      return UInt64
    elseif i == 2
      return Int128
    elseif i == 3 
      return UInt128
    elseif i == 4 
      return Int32
    elseif i == 5 
      return UInt32
    elseif i == 6
      return Int16
    elseif i == 7
      return UInt16
    elseif i == 8 
      return Int8
    elseif i == 9
      return UInt8
    # elseif i == 10
    #   return BigInt
    else
      return Bool
    end
end

function fuzz(f, argtypes...; maxdepth = typemax(Int64))
    worklist = Any[]
    args_ctypes = map(x->Base.filter(supported, concretize(x)), argtypes)
    for types in Iterators.product(args_ctypes...)
        initial_args = map(generate, types)
        push!(worklist, (0, initial_args, nothing))
    end
    fuzz_worklist(f, worklist, maxdepth)
end

function fuzz_wargs(f, initial_args...; maxdepth = typemax(Int64))
    worklist = Any[(0, initial_args, nothing)]
    fuzz_worklist(f, worklist, maxdepth)
end

###
#    Iterative breath first tree search
#      - Invalidated earliest branch and use Z3 to generate an example for the
#        opposite branch
#      - explore all sides and iterate through the program to discover all branches
function fuzz_worklist(f, worklist::Vector{Any}, maxdepth)
    tested = Any[] # did not error
    errored = Any[]

    while !isempty(worklist)
        depth, args, subs = pop!(worklist)
        @debug "Executing item from worklist" depth args subs
        val, _, trace, record = execute(f, args...; subs = subs)

        if val isa Exception
            push!(errored, (val, args, subs))
        else
            push!(tested, (val, args, subs))
        end

        stream = filter(trace)
        branches = Any[]
        for n in (depth+1):nbranch(stream)
            cut_stream = force_and_cut(stream, n)
            push!(branches, (n, cut_stream))
        end

        @info "Found $(length(branches)) new branches to explore"

        for (d, branch) in branches
            if d >= maxdepth
                @info "Terminated fuzzing that went to deep" maxdepth
                continue
            end
            sat, inputs = try
                checkStream(branch)
            catch ex
                @error "Error in Z3 run, skipping" exception=ex branch
                continue
            end
            if sat
                args, subs, _ = inputs
                subs = augment(record, subs)
                # record to subs
                push!(worklist, (d, args, subs))
            end
        end
    end
    tested, errored
end

<|start_filename|>src/tracentaint.jl<|end_filename|>

# A symbol is a name and the concolicly computed type
struct Sym
    name::Symbol
    _type::DataType
    Sym(base, _type::DataType) = new(Base.gensym(Symbol(base, '#', _type)), _type)
    Sym(_type::DataType) = new(Base.gensym(Symbol(_type)), _type)
end
Cassette.metadatatype(::Type{<:TraceCtx}, ::DataType) = Sym

record(args, ctx) = map(args) do arg
    if istagged(arg, ctx)
        metadata(arg, ctx)
    else
        arg
    end
end

# A callsite is a function with a concolic argument set
# children are all the callsites within the function
mutable struct Callsite
    f::Any
    args::Any
    retval::Union{Some{Any}, Nothing}
    children::Vector{Callsite}
    Callsite(f, args) = new(f, args, nothing, Callsite[])
end
Base.push!(trace::Callsite, call::Callsite) = push!(trace.children, call)

struct Metadata
    trace::Callsite
    # For current execution
    record::Vector{Sym}
    substitutes::Dict{DataType, Any}
end
Metadata() = Metadata(
    Callsite(:toplevel, ()), 
    Vector{Sym}(), 
    Dict{DataType, Vector}())

Metadata(m::Metadata, call::Callsite) = Metadata(call, m.record, m.substitutes)
Metadata(subs) = Metadata(Callsite(:toplevel, ()), Vector{Sym}(), subs)

function record!(ctx, type)
    m = ctx.metadata
    sym = Sym(:fval, type)
    push!(m.record, sym)

    if haskey(m.substitutes, type)
        subs = m.substitutes[type]
        if !isempty(subs)
            return pop!(subs), sym
        end
    end
    return nothing, sym
end
Base.push!(m::Metadata, call) = push!(m.trace, call)

function augment(record, subs)
    substitutions = Dict{DataType, Vector}()
    function insert!(type, val)
        if !haskey(substitutions, type)
            substitutions[type] = Vector{type}()
        end
        push!(substitutions[type], val)
    end

    for sym in record
        insert!(sym._type, subs[sym.name])
    end
    return substitutions
end

function Cassette.overdub(ctx::TraceCtx, ::typeof(rand), ::Type{T}) where T<:INTEGERS
    val, sym = record!(ctx, T)
    if val === nothing
        val = rand(T)
    end
    return tag(val::T, ctx, sym)
end

# oftype will drop tag on the x value since we are only looking at it'd type.
function Cassette.overdub(ctx::TraceCtx, ::typeof(oftype), x::Tagged, y)
    x = untag(x, ctx)
    return oftype(x, y)
end

##
# We need to manually override for IntrinsicFunctions (which are the leaf-nodes we are interested in)
# Since in tagged contexts there is an automatic fallback available.
##
function Cassette.overdub(ctx::TraceCtx{Metadata, <:Cassette.Tag}, f::Core.IntrinsicFunction, args...)
    call = Callsite(f, record(args, ctx))
    push!(ctx.metadata, call)
    retval = fallback(ctx, f, args...)
    if any(a -> istagged(a, ctx), args)
        retval = tag(retval, ctx, Sym(typeof(retval)))
    end
    call.retval = Some(istagged(retval, ctx) ? metadata(retval, ctx) : retval)
    return retval
end

##
# Recursivly trace though our program.
#
# Note: This won't trace functions that are defined primitive and there are several
#       fallbacks that Cassette provides for tagged contexts.
#
# Question: Shouldn't canoverdub for primitives be false?
##
function Cassette.overdub(ctx::TraceCtx, f, args...)
    # We need to push the callsite first into metadata
    # otherwise we run into issues with functions that throw errors
    call = Callsite(f, record(args, ctx))
    push!(ctx.metadata, call)

    # Recurse into the next function
    retval = if canrecurse(ctx, f, args...)
        newctx = similarcontext(ctx, metadata = Metadata(ctx.metadata, call))
        Cassette.recurse(newctx, f, args...)
    else
        retval = fallback(ctx, f, args...)
        # If any of our inputs was tagged we want to tag the return value so that
        # the symbolic nature of this operation is captured.
        if any(a -> istagged(a, ctx), args) && !istagged(retval, ctx)
            retval = tag(retval, ctx, Sym(typeof(retval)))
        end
        retval
    end
    call.retval = Some(istagged(retval, ctx) ? metadata(retval, ctx) : retval)
    return retval
end


<|start_filename|>src/asserts.jl<|end_filename|>

"""
    assert(x::Bool)

Assert that the condition generating `x` holds true.
Returns `x`

"""
assert(x::Bool) = x 

"""
    prove(x::Bool)

```julia
assert(x < 10)
y = x - 10
prove(y < 0)
```

Will turn into the corresponding SMT2 program:

```
(declare-const x Int)
(declare-const y Int)
(assert (< x 10))
(assert (= y (- x 10)))
(assert (not (< y 0)))
(check-sat)
```

If the model is not satisfiable, given the pre-condition defined by `assert`
the post-condition given by `prove` holds.

For more information see `check`.
"""
prove(x::Bool) = x

nodub(f, args...) = Expr(:call, Expr(:nooverdub, f), args...)

# Define a Cassette pass to insert asserts after branches
function insertasserts(ctx, ref)
    ir = ref.code_info
    locations = Int[]
    Cassette.insert_statements!(ir.code, ir.codelocs,
        (stmt, i) -> Base.Meta.isexpr(stmt, :gotoifnot) ? 2 : nothing, 
        (stmt, i) -> [
            Expr(:call, ConcolicFuzzer.assert, stmt.args[1]),
            stmt
        ])

    # Taint foreigncalls
    # TODO: llvmcall, check if stmts.args[2] supported
    # This works but is way to aggressive.
    #=
    Cassette.insert_statements!(ir.code, ir.codelocs,
        (stmt, i) -> begin
            stmt = Base.Meta.isexpr(stmt, :(=)) ? stmt.args[2] : stmt
            if Base.Meta.isexpr(stmt, :foreigncall) && supported(stmt.args[2])
                12
            else
                nothing
            end
        end,
        (stmt, i) -> begin
            if Base.Meta.isexpr(stmt, :(=))
                tagstmt = Expr(:(=), stmt.args[1])
                stmt = stmt.args[2]
            else
                tagstmt = nothing
            end
            # setup slot for if-else
            slot = Core.SlotNumber(length(ir.slotflags) + 1)
            push!(ir.slotflags, 0x10) # TODO check flag value
            push!(ir.slotnames, gensym(:temp))
            # obtain location as stacktrace
            type = stmt.args[2]

            # Helper to more easily create result
            result = Any[]
            insert!(stmt) = (id = i+length(result); push!(result, stmt); Core.SSAValue(id))

            #=
            stack = StackTraces.stacktrace()
            loc = StackTraces.remove_frames!(stack, :execute)
            val, sym = record!(ctx, loc, type)
            if val === nothing
                val = stmt
            end
            return tag(val, ctx, sym)
            =#
            stack  = insert!(nodub(StackTraces.stacktrace))
            loc    = insert!(nodub(StackTraces.remove_frames!, stack, QuoteNode(:execute)))
            record = insert!(nodub(ConcolicFuzzer.record!, Expr(:contextslot), loc, type))
            val    = insert!(nodub(Base.getindex, record, 1))
            sym    = insert!(nodub(Base.getindex, record, 2))
            cond   = insert!(nodub(Core.:(===), val, nothing))
            ifstmt = insert!(Expr(:gotoifnot, cond, #=fill in later=#))
            bb1    = insert!(Expr(:(=), slot, stmt))
            elstmt = insert!(Expr(:goto, #=fill in later=#))
            bb2    = insert!(Expr(:(=), slot, val))
            tag    = insert!(nodub(Cassette.tag, slot , Expr(:contextslot), sym))

            # fixup if and else
            push!(result[ifstmt.id - i + 1].args, bb2.id)
            result[elstmt.id - i + 1] = Core.GotoNode(tag.id)

            if tagstmt === nothing
                push!(result, tag)
            else
                push!(tagstmt.args, tag)
                push!(result, tagstmt)
            end
            return result
        end)
    =#
    ir.ssavaluetypes = length(ir.code)
    return ir
end

const InsertAssertsPass = Cassette.@pass insertasserts

# We can't use primitives since otherwise our tracer won't trace assert and prove statements 
# Which is the entire point.
function Cassette.overdub(ctx::TraceCtx, f::F, x::Tagged) where F <: Union{Core.typeof(assert), Core.typeof(prove)}
    call = Callsite(f, record((x, ), ctx))
    retval = untag(x, ctx)
    call.retval = Some(retval)
    push!(ctx.metadata, call)
    return retval 
end