diff --git a/src/csg_priority_enumerator.jl b/src/csg_priority_enumerator.jl
index 00ba93c..82e2803 100644
--- a/src/csg_priority_enumerator.jl
+++ b/src/csg_priority_enumerator.jl
@@ -291,7 +291,7 @@ function _expand(
 )::Union{ExpandFailureReason, Vector{TreeConstraints}}
     nodes::Vector{TreeConstraints} = []
     
-    for rule_index ∈ derivation_heuristic(findall(node.domain), context)
+    for rule_index ∈ derivation_heuristic(findall(node.domain), context, grammar)
         new_node = RuleNode(rule_index, grammar)
 
         # If dealing with the root of the tree, propagate here
diff --git a/src/enumerator_constructors.jl b/src/enumerator_constructors.jl
index e91ad33..0d65f89 100644
--- a/src/enumerator_constructors.jl
+++ b/src/enumerator_constructors.jl
@@ -12,7 +12,7 @@ function get_bfs_enumerator(
     max_size::Int, 
     sym::Symbol, 
     hole_heuristic::Function=heuristic_leftmost, 
-    derivation_heuristic::Function=(a,_) -> a
+    derivation_heuristic::Function=(a,_,_) -> a
 )::ContextSensitivePriorityEnumerator
     expand_function(node, grammar, max_depth, max_holes, context) = _expand(node, grammar, max_depth, max_holes, context, hole_heuristic, derivation_heuristic)
     return ContextSensitivePriorityEnumerator(cfg2csg(grammar), max_depth, max_size, bfs_priority_function, expand_function, sym)
@@ -29,7 +29,7 @@ function get_bfs_enumerator(
     max_size::Int, 
     sym::Symbol,
     hole_heuristic::Function=heuristic_leftmost, 
-    derivation_heuristic::Function=(a,_) -> a
+    derivation_heuristic::Function=(a,_,_) -> a
 )::ContextSensitivePriorityEnumerator
     expand_function(node, grammar, max_depth, max_holes, context) = _expand(node, grammar, max_depth, max_holes, context, hole_heuristic, derivation_heuristic)
     return ContextSensitivePriorityEnumerator(grammar, max_depth, max_size, bfs_priority_function, expand_function, sym)
@@ -49,7 +49,7 @@ function get_dfs_enumerator(
     max_size::Int, 
     sym::Symbol,
     hole_heuristic::Function=heuristic_leftmost, 
-    derivation_heuristic::Function=(a,_) -> a
+    derivation_heuristic::Function=(a,_,_) -> a
 )::ContextSensitivePriorityEnumerator
     expand_function(node, grammar, max_depth, max_holes, context) = _expand(node, grammar, max_depth, max_holes, context, hole_heuristic, derivation_heuristic)
     return ContextSensitivePriorityEnumerator(cfg2csg(grammar), max_depth, max_size, dfs_priority_function, expand_function, sym)
@@ -66,7 +66,7 @@ function get_dfs_enumerator(
         max_size::Int, 
         sym::Symbol,
         hole_heuristic::Function=heuristic_leftmost, 
-        derivation_heuristic::Function=(a,_) -> a    
+        derivation_heuristic::Function=(a,_,_) -> a    
 )::ContextSensitivePriorityEnumerator
     expand_function(node, grammar, max_depth, max_holes, context) = _expand(node, grammar, max_depth, max_holes, context, hole_heuristic, derivation_heuristic)
     return ContextSensitivePriorityEnumerator(grammar, max_depth, max_size, dfs_priority_function, expand_function, sym)
@@ -97,7 +97,7 @@ function get_most_likely_first_enumerator(
     max_size::Int, 
     sym::Symbol,
     hole_heuristic::Function=heuristic_leftmost, 
-    derivation_heuristic::Function=(a,_) -> a
+    derivation_heuristic::Function=(a,_,_) -> a
 )::ContextSensitivePriorityEnumerator
     expand_function(node, grammar, max_depth, max_holes, context) = _expand(node, grammar, max_depth, max_holes, context, hole_heuristic, derivation_heuristic)
     return ContextSensitivePriorityEnumerator(cfg2csg(grammar), max_depth, max_size, most_likely_priority_function, expand_function, sym)
@@ -115,7 +115,7 @@ function get_most_likely_first_enumerator(
     max_size::Int, 
     sym::Symbol,
     hole_heuristic::Function=heuristic_leftmost, 
-    derivation_heuristic::Function=(a,_) -> a
+    derivation_heuristic::Function=(a,_,_) -> a
 )::ContextSensitivePriorityEnumerator
     expand_function(node, grammar, max_depth, max_holes, context) = _expand(node, grammar, max_depth, max_holes, context, hole_heuristic, derivation_heuristic)
     return ContextSensitivePriorityEnumerator(grammar, max_depth, max_size, most_likely_priority_function, expand_function, sym)