Adding conflict analysis extenstion to HerbSearch

Project.toml

@@ -17,9 +17,11 @@ StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 
 [weakdeps]
+Satisfiability = "160ab843-0bc6-4ba4-9585-b7478b70f443"
 DecisionTree = "7806a523-6efd-50cb-b5f6-3fa6f1930dbb"
 
 [extensions]
+ConflictAnalysisExt = "Satisfiability"
 DivideAndConquerExt = "DecisionTree"
 
 [compat]
@@ -33,6 +35,7 @@ HerbInterpret = "0.2.1"
 HerbSpecification = "0.2"
 MLStyle = "0.4.17"
 Random = "1.8.0"
+Satisfiability = "0.2.0"
 StatsBase = "0.34"
 TimerOutputs = "0.5.28"
 julia = "1.10"

ext/ConflictAnalysisExt/ConflictAnalysisExt.jl

@@ -0,0 +1,191 @@
+module ConflictAnalysisExt
+
+using HerbSearch
+using HerbCore
+using HerbGrammar
+using HerbSpecification
+using HerbConstraints
+using HerbInterpret
+using DataStructures
+using Satisfiability
+using DocStringExtensions
+using MLStyle
+
+
+include("data.jl")
+include("techniques/muc.jl")
+include("techniques/era.jl")
+include("techniques/sean.jl")
+include("pipeline.jl")
+
+
+"""
+		$(TYPEDSIGNATURES)
+Synthesizes a program using conflict analysis. 
+
+Uses the defined techniques, implementations can be found in the techniques folder, and synthesizes a program while eliminating candidate programs during search. 
+
+# Arguments 
+- `problem::Problem` : Specification of the program synthesis problem.
+- `iterator::ProgramIterator` : Iterator over candidate programs that is used to search for solutions of the sub-programs.
+- `interpret::Function` : Function for custom interpreting a candidate solution and input.
+- `max_time::Int` : Maximum time that the iterator will run 
+- `max_enumerations::Int` : Maximum number of iterations that the iterator will run 
+- `mod::Module` : A module containing definitions for the functions in the grammar. Defaults to `Main`.
+- `techniques::Vector{Symbol}` : A vector of symbols representing the conflict analysis techniques to use.
+
+Returns the `RuleNode` representing the final program constructed from the solutions to the subproblems. Can also return `nothing` if no solution is found within the constraints.
+"""
+function HerbSearch.conflict_analysis(
+    problem::Problem,
+    iterator::ProgramIterator,
+    interpret::Union{Function, Nothing} = nothing;
+    max_time = typemax(Int),
+    max_enumerations = typemax(Int),
+    mod::Module = Main,
+    techniques::Vector{Symbol} = [:ERA, :MUC, :SeAn]
+
+)::Tuple{Union{AbstractRuleNode, Nothing}, Int64, Int64}
+    start_time   = time()
+    solver       = iterator.solver
+    grammar      = get_grammar(solver)
+    grammar_tags = isnothing(interpret) ? nothing : get_relevant_tags(grammar)
+    symboltable  = grammar2symboltable(grammar, mod)
+    counter      = 0
+    cons_counter = 0
+
+    techs = build_techniques(techniques)
+
+    for (i, candidate_program) ∈ enumerate(iterator)
+        counter = i
+        expr = rulenode2expr(candidate_program, grammar)
+        output, result, counter_example = isnothing(interpret) ? 
+            evaluate(expr, problem, symboltable) : 
+            evaluate(candidate_program, problem, grammar_tags, interpret)
+
+        if result == success
+            return (freeze_state(candidate_program), counter, cons_counter)
+        else
+            ctx = ConflictContext(grammar, symboltable, candidate_program, output, counter_example)
+            constraints, grammar_constraints = run_conflict_pipeline(techs, ctx)
+
+            for c in grammar_constraints
+                addconstraint!(grammar, c.cons)
+            end
+            if !isempty(constraints)
+                HerbSearch.add_constraints!(iterator, AbstractGrammarConstraint[c.cons for c in constraints])
+            end
+
+            cons_counter += length(constraints) + length(grammar_constraints)
+        end
+
+        if i > max_enumerations || time() - start_time > max_time
+            println("Stopping criteria met")
+            break
+        end
+    end
+
+    # Clean up
+    for t in techs
+        try
+            close(t)
+        catch _
+            # ignore if technique has no resources
+        end
+    end
+
+    return (nothing, counter, cons_counter)
+end
+
+"""
+Gets relevant symbol to easily match grammar rules to operations in `interpret` function
+"""
+function get_relevant_tags(grammar::ContextSensitiveGrammar)
+    tags = Dict{Int,Any}()
+    for (ind, r) in pairs(grammar.rules)
+        tags[ind] = if typeof(r) != Expr
+            r
+        else
+            @match r.head begin
+                :block => :OpSeq
+                :call => r.args[1]
+                :if => :IF
+            end
+        end
+    end
+    return tags
+end
+
+"""
+    execute_on_input(tab::SymbolTable, expr::Any, input::Dict{Symbol, T}, interpret::Function)::Any where T
+
+Custom execute_on_input function that uses a given interpret function.
+"""
+function HerbSearch.execute_on_input(program::AbstractRuleNode, grammar_tags::Dict{Int, Any}, input::Dict{Symbol, T}, interpret::Function)::Any where T
+    return interpret(program, grammar_tags, input)
+end
+
+@enum EvalResult success=1 failed=2 crashed=3
+"""
+    evaluate(
+        expr::Any,
+        problem::Problem{<:AbstractVector{<:IOExample}},
+        symboltable::SymbolTable
+    )::Tuple{Union{Any, Nothing}, EvalResult, Union{<:IOExample, Nothing}}
+
+Evaluate the expression on the examples using the given symboltable.
+"""
+function evaluate(
+    expr::Any,
+    problem::Problem{<:AbstractVector{<:IOExample}},
+    symboltable::SymbolTable
+)::Tuple{Union{Any, Nothing}, EvalResult, Union{<:IOExample, Nothing}}
+    output = nothing
+
+    for example ∈ problem.spec
+        try
+            output = execute_on_input(symboltable, expr, example.in)
+            if (output != example.out)
+                return (output, failed, example)
+            end
+        catch e
+            return (nothing, crashed, example)
+        end
+    end
+
+    return (output, success, nothing)
+end
+
+"""
+    evaluate(
+        program::AbstractRuleNode,
+        problem::Problem{<:AbstractVector{<:IOExample}},
+        grammar_tags::Dict{Int, Any},
+        interpret::Union{Function, Nothing} = nothing
+    )::Tuple{Union{Any, Nothing}, EvalResult, Union{<:IOExample, Nothing}}
+
+Evaluate the program on the examples using a custom interpret function if provided.
+"""
+function evaluate(
+    program::AbstractRuleNode,
+    problem::Problem{<:AbstractVector{<:IOExample}},
+    grammar_tags::Dict{Int, Any},
+    interpret::Union{Function, Nothing} = nothing
+)::Tuple{Union{Any, Nothing}, EvalResult, Union{<:IOExample, Nothing}}
+    output = nothing
+
+    for example ∈ problem.spec
+        try
+            output = execute_on_input(program, grammar_tags, example.in, interpret)
+            if (output != example.out)
+                return (output, failed, example)
+            end
+        catch e
+            return (nothing, crashed, example)
+        end
+    end
+
+    return (output, success, nothing)
+end
+
+end # module

ext/ConflictAnalysisExt/data.jl

@@ -0,0 +1,76 @@
+abstract type AbstractConflictTechnique end
+
+"Generic input data for conflict analysis, used by check_conflict"
+abstract type AbstractConflictInput end
+
+struct MUCInput <: AbstractConflictInput
+    root::AbstractRuleNode
+    grammar::AbstractGrammar
+    counter_example::IOExample
+end
+
+struct ERAInput <: AbstractConflictInput
+    root::AbstractRuleNode
+    grammar::AbstractGrammar
+    evaluation::Union{Any, Nothing}
+end
+
+struct SeAnInput <: AbstractConflictInput
+    root::AbstractRuleNode
+    grammar::AbstractGrammar
+    symboltable::SymbolTable
+    counter_example::IOExample
+end
+
+"Generic conflict data returned by check_conflict, used by analyse_conflict"
+abstract type AbstractConflictData end
+
+struct UnsatCoreData <: AbstractConflictData
+    core_constraints_map::Dict{Int, Vector{Tuple{Any, String, Any}}}
+end
+
+struct ERAException <: AbstractConflictData
+    program::AbstractRuleNode
+end
+
+struct ERARedundantValues <: AbstractConflictData
+    program::AbstractRuleNode
+    value::Any
+end
+
+struct SeAnData <: AbstractConflictData
+    semantics::Vector{Expr}
+end
+
+"Generic conflict constraint returned by analyse_conflict"
+abstract type AbstractConflictConstraint end
+
+struct MUCConstraint <: AbstractConflictConstraint
+    cons::Forbidden
+    origin::String
+    add_to_grammar::Bool
+end
+
+function MUCConstraint(cons::Forbidden, origin::String)
+    return MUCConstraint(cons, origin, false)
+end
+
+struct ERAConstraint <: AbstractConflictConstraint
+    cons::Forbidden
+    origin::String
+    add_to_grammar::Bool
+end
+
+function ERAConstraint(cons::Forbidden, origin::String)
+    return ERAConstraint(cons, origin, true)
+end
+
+struct SeAnConstraint <: AbstractConflictConstraint
+    cons::Forbidden
+    origin::String
+    add_to_grammar::Bool
+end
+
+function SeAnConstraint(cons::Forbidden, origin::String)
+    return SeAnConstraint(cons, origin, false)
+end

ext/ConflictAnalysisExt/pipeline.jl

@@ -0,0 +1,77 @@
+struct ConflictContext
+    grammar::AbstractGrammar
+    symboltable::SymbolTable
+    candidate::AbstractRuleNode
+    output::Any
+    counter_example::IOExample
+end
+
+"""
+    run_conflict_pipeline(techniques::Vector{<:AbstractConflictTechnique})
+
+Run each conflict analysis technique in parallel by first calling `check_conflict`,
+and if a conflict is found, calls `analyze_conflict`. Collects and returns all
+resulting `AbstractConflictConstraint` objects in a thread-safe manner.
+"""
+function run_conflict_pipeline(techniques::Vector{<:AbstractConflictTechnique}, ctx::ConflictContext)
+    constraints = AbstractConflictConstraint[]
+    grammar_constraints = AbstractConflictConstraint[]
+
+    for tech in techniques
+        autoinput!(tech, ctx)
+
+        tech.data = check_conflict(tech)
+        if isnothing(tech.data) 
+            continue
+        end
+
+        analyze_result = analyze_conflict(tech)
+        if isnothing(analyze_result)
+            continue
+        end
+
+        if analyze_result isa AbstractConflictConstraint
+            if analyze_result.add_to_grammar
+                push!(grammar_constraints, analyze_result)
+            else
+                push!(constraints, analyze_result)
+            end
+        elseif analyze_result isa AbstractVector
+            for c in analyze_result
+                if c isa AbstractConflictConstraint
+                    if c.add_to_grammar 
+                        push!(grammar_constraints, c)
+                    else
+                        push!(constraints, c)
+                    end
+                end
+            end
+        else
+            # Consider logging this or handling more gracefully
+            @error "analyze_conflict returned unsupported type: $(typeof(analyze_result))"
+        end
+    end
+    return constraints, grammar_constraints
+end
+
+
+"""
+    build_techniques(names::Vector{Symbol}) -> Vector{AbstractConflictTechnique}
+
+Instantiate techniques once, based on a simple list of symbols.
+"""
+function build_techniques(names::Vector{Symbol})
+    techs = AbstractConflictTechnique[]
+    for nm in names
+        nm === :MUC  && push!(techs, MUC())
+        nm === :ERA  && push!(techs, ERA())
+        nm === :SeAn && push!(techs, SeAn())
+    end
+    return techs
+end
+
+autoinput!(tech::MUC,  ctx::ConflictContext) = (tech.input = MUCInput(ctx.candidate, ctx.grammar, ctx.counter_example))
+autoinput!(tech::ERA,  ctx::ConflictContext) = (tech.input = ERAInput(ctx.candidate, ctx.grammar, ctx.output))
+autoinput!(tech::SeAn, ctx::ConflictContext) = (tech.input = SeAnInput(ctx.candidate, ctx.grammar, ctx.symboltable, ctx.counter_example))
+
+autoinput!(::AbstractConflictTechnique, ::ConflictContext) = nothing