Probe Implementation

Project.toml

@@ -1,6 +1,10 @@
 name = "Garden"
 uuid = "22d0bc84-1826-4aaf-b584-c2a6a91114a2"
-version = "0.1.0"
+version = "0.1.1"
+
+[deps]
+DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
+Herb = "c09c6b7f-4f63-49de-90d9-97a3563c0f4a"
 
 [deps]
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"

run.jl

@@ -0,0 +1,23 @@
+using Garden
+using Herb
+
+
+grammar = @cfgrammar begin
+    Start = Int
+    Int = Int + Int
+    Int = |(1:5)
+    Int = x
+end
+# problem = Problem( [IOExample{Symbol, Any}(Dict(), 2)])
+problem = Problem([[IOExample(Dict(:x => x), 2x+1) for x ∈ 1:5];
+                   [IOExample(Dict(:x => 2), 0)]])
+
+result = probe(
+    grammar,
+    :Start,
+    problem;
+    max_depth = 4
+)
+
+println(result)
+

src/Garden.jl

@@ -2,8 +2,14 @@ module Garden
 
 using DocStringExtensions
 
+include("utils.jl")
+include("probe/method.jl")
 include("frangel/method.jl")
 
-export FrAngel
+export 
+    Probe,
+    NoProgramFoundError,
+    SynthResult,
+    FrAngel
 
 end # module Garden

src/probe/README.md

@@ -0,0 +1,17 @@
+# Probe
+
+[Publication (Open Access)](https://doi.org/10.1145/3428295)
+
+```
+@article{DBLP:journals/pacmpl/BarkePP20,
+  author       = {Shraddha Barke and
+                  Hila Peleg and
+                  Nadia Polikarpova},
+  title        = {Just-in-time learning for bottom-up enumerative synthesis},
+  journal      = {Proc. {ACM} Program. Lang.},
+  volume       = {4},
+  number       = {{OOPSLA}},
+  pages        = {227:1--227:29},
+  year         = {2020}
+}
+```

src/probe/method.jl

@@ -0,0 +1,160 @@
+module Probe
+
+using DocStringExtensions
+using Garden: SynthResult, optimal_program, suboptimal_program
+using Herb.HerbCore: AbstractRuleNode, AbstractGrammar, rulesoftype
+using Herb.HerbGrammar: normalize!, init_probabilities!, ContextSensitiveGrammar, rulenode2expr, grammar2symboltable
+using Herb.HerbSpecification: AbstractSpecification, Problem
+using Herb.HerbInterpret: SymbolTable
+using Herb.HerbConstraints: freeze_state, get_grammar
+using Herb.HerbSearch: MLFSIterator, evaluate, ProgramIterator, log_probability
+
+
+"""
+    $(TYPEDSIGNATURES)
+
+Synthesize a program using the `grammar` that follows the `spec` following the method from 
+["Just-in-time learning for bottom-up enumerative synthesis"](https://doi.org/10.1145/3428295).
+```
+"""
+function probe(
+        grammar::AbstractGrammar,
+        starting_sym::Symbol,
+        problem::Problem;
+        probe_cycles::Int = 3,
+        max_iterations::Int = typemax(Int),
+        max_iteration_time::Int = typemax(Int),
+        kwargs...
+)::Union{AbstractRuleNode, Nothing}
+    if isnothing(grammar.log_probabilities)
+        init_probabilities!(grammar)
+    end
+
+    for _ in 1:probe_cycles
+        # Gets an iterator with some limit (the low-level budget)
+        iterator = MLFSIterator(grammar, starting_sym; kwargs...)
+
+        # Run a budgeted search
+        promising_programs, result_flag = get_promising_programs_with_fitness(
+            iterator, problem; max_time = max_iteration_time,
+            max_enumerations = max_iterations)
+
+        if result_flag == optimal_program
+            program, score = only(promising_programs) # returns the only element
+            return program
+        end
+
+        # Throw an error if no programs were found.
+        if length(promising_programs) == 0
+            throw(NoProgramFoundError("No promising program found for the given specification. Try exploring more programs."))
+        end
+
+        # Update grammar probabilities 
+        modify_grammar_probe!(promising_programs, grammar)
+    end
+
+    @warn "No solution found within $probe_cycles Probe iterations."
+    return nothing
+end
+
+
+"""
+    $(TYPEDSIGNATURES)
+
+Decide whether to keep a program, or discard it, based on the specification. 
+Returns the portion of solved examples.
+"""
+function decide_probe(
+        program::AbstractRuleNode,
+        problem::Problem,
+        grammar::ContextSensitiveGrammar,
+        symboltable::SymbolTable)::Real
+    expr = rulenode2expr(program, grammar)
+    println("expr:", expr)
+    fitness = evaluate(problem, expr, symboltable, shortcircuit = false)
+    return fitness
+end
+
+"""
+    $(TYPEDSIGNATURES)
+
+Modify the grammar based on the programs kept during the `decide` step.
+Takes a set of programs and their fitnesses, which describe how useful the respective program is.
+Updates a rules probability based on the highest program fitness the rule occurred in. 
+The update function is taken from the Probe paper. Instead of introducing a normalization value, we just call `normalize!` instead.
+"""
+function modify_grammar_probe!(
+        saved_program_fitness::Set{Tuple{<:AbstractRuleNode, Real}},
+        grammar::AbstractGrammar
+)::AbstractGrammar 
+    orig_probs = exp.(grammar.log_probabilities)
+
+    for i in 1:length(grammar.log_probabilities)
+        max_fitness = 0
+
+        # Find maximum fitness for programs with that rule among saved programs
+        for (program, fitness) in saved_program_fitness
+            if !isempty(rulesoftype(program, Set(i))) && fitness > max_fitness
+                max_fitness = fitness                
+            end
+        end
+
+        # Update the probability according to Probe's formula
+        prob = log_probability(grammar, i)
+        orig_probs[i] = log(exp(prob)^(1-max_fitness))
+    end
+    # Normalize probabilities after the update
+    normalize!(grammar)
+
+    return grammar
+end
+
+
+"""
+    $(TYPEDSIGNATURES)
+
+Iterates over the solutions to find partial or full solutions.
+Takes an iterator to enumerate programs. Quits when `max_time` or `max_enumerations` is reached.
+If the program solves the problem, it is returned with the `optimal_program` flag.
+If a program solves some of the problem (e.g. some but not all examples) it is added to the list of `promising_programs`.
+The set of promising programs is returned eventually.
+"""
+function get_promising_programs_with_fitness(
+        iterator::ProgramIterator,
+        problem::Problem;
+        max_time = typemax(Int),
+        max_enumerations = typemax(Int),
+        mod::Module = Main
+)::Tuple{Set{Tuple{AbstractRuleNode, Real}}, SynthResult}
+    start_time = time()
+    grammar = get_grammar(iterator.solver)
+    symboltable::SymbolTable = grammar2symboltable(grammar, mod)
+
+    promising_programs = Set{Tuple{AbstractRuleNode, Real}}()
+
+    for (i, candidate_program) in enumerate(iterator)
+        fitness = decide_probe(candidate_program, problem, grammar, symboltable)
+
+        if fitness == 1
+            push!(promising_programs, (freeze_state(candidate_program), fitness))
+            return (promising_programs, optimal_program)
+        elseif fitness > 0
+            push!(promising_programs, (freeze_state(candidate_program), fitness))
+        end
+
+        # Check stopping criteria
+        if i > max_enumerations || time() - start_time > max_time
+            break
+        end
+    end
+
+    return (promising_programs, suboptimal_program)
+end
+
+export 
+    probe, 
+    decide_probe, 
+    modify_grammar_probe!,
+    get_promising_programs_with_fitness
+
+end

src/probe/ref.bib

@@ -0,0 +1,11 @@
+@article{DBLP:journals/pacmpl/BarkePP20,
+  author       = {Shraddha Barke and
+                  Hila Peleg and
+                  Nadia Polikarpova},
+  title        = {Just-in-time learning for bottom-up enumerative synthesis},
+  journal      = {Proc. {ACM} Program. Lang.},
+  volume       = {4},
+  number       = {{OOPSLA}},
+  pages        = {227:1--227:29},
+  year         = {2020}
+}

src/utils.jl

@@ -0,0 +1,5 @@
+@enum SynthResult optimal_program=1 suboptimal_program=2
+
+struct NoProgramFoundError <: Exception
+    message::String
+end

test/Project.toml

@@ -2,4 +2,4 @@
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 Herb = "c09c6b7f-4f63-49de-90d9-97a3563c0f4a"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
-Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

test/test_probe.jl

@@ -0,0 +1,56 @@
+using Herb.HerbCore: @rulenode, RuleNode
+using Herb.HerbGrammar: @cfgrammar
+using Herb.HerbSpecification: IOExample, Problem
+using Garden: Probe, NoProgramFoundError, SynthResult
+using .Probe: probe, get_promising_programs_with_fitness, modify_grammar_probe!
+
+@testset "Probe" begin
+    @testset verbose=true "Integration tests" begin
+        # Define extra grammar as FrAngel will change it.
+        grammar = @cfgrammar begin
+            Start = Int
+            Int = Int + Int
+            Int = |(1:5)
+        end
+
+        problem = Problem([IOExample(Dict(:x => x), 2x+1) for x ∈ 1:5])
+        result = probe(
+            grammar,
+            :Start,
+            problem;
+            max_depth = 4
+        )
+
+        @test rulenode2expr(result, grammar) == 2
+
+        imp_problem = Problem([[IOExample(Dict(:x => x), 2x+1) for x ∈ 1:5];
+                   [IOExample(Dict(:x => 2), 0)]])
+
+        # A program yielding 0 is impossible to derive from the grammar.
+        @test_throws NoProgramFoundError probe(
+            grammar, :Start, imp_problem; max_depth = 3)
+    end
+
+    grammar = @cfgrammar begin
+        Start = Int
+        Int = Int + Int
+        Int = |(1:5)
+        Int = x
+    end
+
+    @testset "modify_grammar_probe!" begin
+        program = @rulenode 2{3, 4}
+        fitness = 1
+
+        orig_probs = grammar.log_probabilities
+
+        modify_grammar_probe!(Set{Tuple{RuleNode, Real}}((program, fitness)), grammar)
+
+        new_probs = grammar.log_probabilities
+        # Probabilities change
+        @test orig_probs != grammar.log_probabilities
+        # Test increase
+        @test maximum(new_probs[[2,3,4]]) < minimum(orig_probs)
+        @test minimum(new_probs[[1,5,6,7,8]]) > maximum(orig_probs)
+    end
+end