Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

A bit cleanup; add type for env #143

Merged
merged 1 commit into from
Feb 27, 2024
Merged

A bit cleanup; add type for env #143

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
54 changes: 25 additions & 29 deletions src/compiler_pass.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,48 +1,44 @@
abstract type CompilerPass end

function analyze_program(pass::CompilerPass, expr::Expr, env, vargs...)
for ex in expr.args
if Meta.isexpr(ex, :(=))
analyze_assignment(pass, ex, env, vargs...)
elseif Meta.isexpr(ex, :call) && ex.args[1] == :(~)
analyze_assignment(pass, ex, env, vargs...)
elseif Meta.isexpr(ex, :for)
analyze_for_loop(pass, ex, env, vargs...)
@inline is_deterministic(expr::Expr) = Meta.isexpr(expr, :(=))
@inline is_stochastic(expr::Expr) = Meta.isexpr(expr, :call) && expr.args[1] == :(~)

function analyze_program(pass::CompilerPass, expr::Expr, env::NamedTuple)
for statement in expr.args
if is_deterministic(statement) || is_stochastic(statement)
analyze_assignment(pass, statement, env)
elseif Meta.isexpr(statement, :for)
analyze_for_loop(pass, statement, env)
else
error("Unsupported expression in top level: $ex")
error("Unsupported expression in top level: $statement")
end
end
return post_process(pass, expr, env, vargs...)
return post_process(pass, expr, env)
end

function analyze_for_loop(pass::CompilerPass, expr, env, vargs...)
function analyze_for_loop(pass::CompilerPass, expr::Expr, env::NamedTuple)
loop_var, lb, ub, body = decompose_for_expr(expr)
lb = Int(evaluate(lb, env))
ub = Int(evaluate(ub, env))

for i in lb:ub
for ex in body.args
if Meta.isexpr(ex, :(=))
analyze_assignment(
pass, ex, merge(env, NamedTuple{(loop_var,)}((i,))), vargs...
)
elseif Meta.isexpr(ex, :call) && ex.args[1] == :(~)
analyze_assignment(
pass, ex, merge(env, NamedTuple{(loop_var,)}((i,))), vargs...
)
elseif Meta.isexpr(ex, :for)
analyze_for_loop(
pass, ex, merge(env, NamedTuple{(loop_var,)}((i,))), vargs...
)
for statement in body.args
env = merge(env, NamedTuple{(loop_var,)}((i,)))
if is_deterministic(statement) || is_stochastic(statement)
analyze_assignment(pass, statement, env)
elseif Meta.isexpr(statement, :for)
analyze_for_loop(pass, statement, env)
else
error("Unsupported expression in for loop body: $ex")
error("Unsupported expression in for loop body: $statement")
end
end
end
end

function analyze_assignment end

function post_process end

@enum VariableTypes begin
Logical
Stochastic
Expand Down Expand Up @@ -327,7 +323,7 @@ is_resolved(::Array{Missing}) = false
is_resolved(::Union{Symbol,Expr}) = false
is_resolved(::Any) = false

function analyze_assignment(pass::CollectVariables, expr::Expr, env)
function analyze_assignment(pass::CollectVariables, expr::Expr, env::NamedTuple)
if Meta.isexpr(expr, :(=))
lhs_expr = expr.args[1]
else # Expr(:call, :(~), ...)
Expand Down Expand Up @@ -416,7 +412,7 @@ function has_value(transformed_variables, v::Var)
end
end

function analyze_assignment(pass::ConstantPropagation, expr::Expr, env)
function analyze_assignment(pass::ConstantPropagation, expr::Expr, env::NamedTuple)
if Meta.isexpr(expr, :(=)) && !should_skip_eval(expr.args[2])
lhs = find_variables_on_lhs(expr.args[1], env)

Expand Down Expand Up @@ -480,7 +476,7 @@ function PostChecking(data, transformed_variables::Dict)
)
end

function analyze_assignment(pass::PostChecking, expr::Expr, env)
function analyze_assignment(pass::PostChecking, expr::Expr, env::NamedTuple)
@inline set_value!(d::Dict, value, v::Scalar) = d[v.name] = value
@inline set_value!(d::Dict, value, v::Var) = d[v.name][v.indices...] = value
@inline get_value(d::Dict, v::Scalar) = d[v.name]
Expand Down Expand Up @@ -806,7 +802,7 @@ try_cast_to_int(x::Integer) = x
try_cast_to_int(x::Real) = Int(x) # will error if !isinteger(x)
try_cast_to_int(x) = x # catch other types, e.g. UnitRange, Colon

function analyze_assignment(pass::NodeFunctions, expr::Expr, env)
function analyze_assignment(pass::NodeFunctions, expr::Expr, env::NamedTuple)
@capture(expr, lhs_expr_ ~ rhs_expr_) || @capture(expr, lhs_expr_ = rhs_expr_)
var_type = Meta.isexpr(expr, :(=)) ? Logical : Stochastic

Expand Down
Loading