<commit_before>package mb.statix.random.strategy;

import static mb.statix.constraints.Constraints.conjoin;

import java.util.Iterator;
import java.util.List;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import org.apache.commons.math3.distribution.EnumeratedDistribution;
import org.apache.commons.math3.util.Pair;

import com.google.common.collect.Lists;
import com.google.common.collect.Sets;

import io.usethesource.capsule.Map;
import io.usethesource.capsule.Set;
import mb.nabl2.terms.unification.IUnifier;
import mb.nabl2.util.ImmutableTuple3;
import mb.nabl2.util.Tuple3;
import mb.statix.constraints.CConj;
import mb.statix.constraints.CInequal;
import mb.statix.constraints.CTrue;
import mb.statix.constraints.CUser;
import mb.statix.random.FocusedSearchState;
import mb.statix.random.SearchContext;
import mb.statix.random.SearchState;
import mb.statix.random.SearchStrategy;
import mb.statix.random.nodes.SearchNode;
import mb.statix.random.nodes.SearchNodes;
import mb.statix.random.util.RandomGenerator;
import mb.statix.random.util.StreamUtil;
import mb.statix.scopegraph.reference.CriticalEdge;
import mb.statix.solver.Delay;
import mb.statix.solver.IConstraint;
import mb.statix.solver.IState;
import mb.statix.solver.completeness.ICompleteness;
import mb.statix.spec.ApplyResult;
import mb.statix.spec.Rule;
import mb.statix.spec.RuleUtil;

final class Expand extends SearchStrategy<FocusedSearchState<CUser>, SearchState> {
    private final double defaultWeight;
    private final java.util.Map<String, Double> weights;

    Expand(double defaultWeight, java.util.Map<String, Double> weights) {
        this.defaultWeight = defaultWeight;
        this.weights = weights;
    }

    @Override protected SearchNodes<SearchState> doApply(SearchContext ctx,
            SearchNode<FocusedSearchState<CUser>> node) {
        final FocusedSearchState<CUser> input = node.output();
        final CUser predicate = input.focus();

        final List<Rule> rules = input.state().spec().rules().get(predicate.name());
        final java.util.Map<String, Long> ruleCounts =
                rules.stream().collect(Collectors.groupingBy(Rule::label, Collectors.counting()));

        // next block is a almost a copy of the CUser case in the solver and must be kept in sync
        final Iterator<Rule> it = rules.iterator();
        final Set.Transient<Tuple3<Rule, ApplyResult, IConstraint>> _results = Set.Transient.of();
        IConstraint unguard = new CTrue(predicate); // we build this iteratively for all previous rules we ignored
        while(it.hasNext()) {
            final Rule rule = it.next();
            final ApplyResult applyResult;
            if((applyResult = RuleUtil.apply(input.state(), rule, predicate.args(), predicate).orElse(null)) == null) {
                // ignore
            } else if(applyResult.guard().isEmpty()) {
                if(!_results.isEmpty()) {
                    throw new IllegalStateException("Rule order must be wrong");
                }
                _results.__insert(ImmutableTuple3.of(rule, applyResult, new CTrue(predicate)));
                break;
            } else {
                final IConstraint _unguard = conjoin(applyResult.guard().entrySet().stream()
                        .map(e -> new CInequal(e.getKey(), e.getValue(), predicate)).collect(Collectors.toList()));
                _results.__insert(ImmutableTuple3.of(rule, applyResult, unguard));
                unguard = new CConj(unguard, _unguard, predicate);
            }
        }
        final Set.Immutable<Tuple3<Rule, ApplyResult, IConstraint>> results = _results.freeze();

        final List<Pair<SearchNode<SearchState>, Double>> newNodes = Lists.newArrayList();<commit_msg>Fix to logic in rule selection.<commit_after>package mb.statix.random.strategy;

import static mb.statix.constraints.Constraints.conjoin;

import java.util.Iterator;
import java.util.List;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import org.apache.commons.math3.distribution.EnumeratedDistribution;
import org.apache.commons.math3.util.Pair;

import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Lists;
import com.google.common.collect.Sets;

import io.usethesource.capsule.Map;
import io.usethesource.capsule.Set;
import mb.nabl2.terms.ITerm;
import mb.nabl2.terms.ITermVar;
import mb.nabl2.terms.unification.IUnifier;
import mb.nabl2.util.ImmutableTuple3;
import mb.nabl2.util.Tuple3;
import mb.statix.constraints.CConj;
import mb.statix.constraints.CInequal;
import mb.statix.constraints.CTrue;
import mb.statix.constraints.CUser;
import mb.statix.random.FocusedSearchState;
import mb.statix.random.SearchContext;
import mb.statix.random.SearchState;
import mb.statix.random.SearchStrategy;
import mb.statix.random.nodes.SearchNode;
import mb.statix.random.nodes.SearchNodes;
import mb.statix.random.util.RandomGenerator;
import mb.statix.random.util.StreamUtil;
import mb.statix.scopegraph.reference.CriticalEdge;
import mb.statix.solver.Delay;
import mb.statix.solver.IConstraint;
import mb.statix.solver.IState;
import mb.statix.solver.completeness.ICompleteness;
import mb.statix.spec.ApplyResult;
import mb.statix.spec.Rule;
import mb.statix.spec.RuleUtil;

final class Expand extends SearchStrategy<FocusedSearchState<CUser>, SearchState> {
    private final double defaultWeight;
    private final java.util.Map<String, Double> weights;

    Expand(double defaultWeight, java.util.Map<String, Double> weights) {
        this.defaultWeight = defaultWeight;
        this.weights = weights;
    }

    @Override protected SearchNodes<SearchState> doApply(SearchContext ctx,
            SearchNode<FocusedSearchState<CUser>> node) {
        final FocusedSearchState<CUser> input = node.output();
        final CUser predicate = input.focus();

        final List<Rule> rules = input.state().spec().rules().get(predicate.name());
        final java.util.Map<String, Long> ruleCounts =
                rules.stream().collect(Collectors.groupingBy(Rule::label, Collectors.counting()));

        // next block is a almost a copy of the CUser case in the solver and must be kept in sync
        final Iterator<Rule> it = rules.iterator();
        final Set.Transient<Tuple3<Rule, ApplyResult, IConstraint>> _results = Set.Transient.of();
        IConstraint unguard = new CTrue(predicate); // we build this iteratively for all previous rules we ignored
        while(it.hasNext()) {
            final Rule rule = it.next();
            final ApplyResult applyResult;
            if((applyResult = RuleUtil.apply(input.state(), rule, predicate.args(), predicate).orElse(null)) == null) {
                // ignore
            } else {
                _results.__insert(ImmutableTuple3.of(rule, applyResult, unguard));
                final ImmutableMap<ITermVar, ITerm> guard = applyResult.guard();
                if(guard.isEmpty()) {
                    break;
                } else {
                    final IConstraint _unguard = conjoin(guard.entrySet().stream()
                            .map(e -> new CInequal(e.getKey(), e.getValue(), predicate)).collect(Collectors.toList()));
                    unguard = new CConj(unguard, _unguard, predicate);
                }
            }
        }
        final Set.Immutable<Tuple3<Rule, ApplyResult, IConstraint>> results = _results.freeze();

        final List<Pair<SearchNode<SearchState>, Double>> newNodes = Lists.newArrayList();
{'fraction_non_alphanumeric': 0.11256335022672713, 'fraction_numerical': 0.003067484662576687, 'mean_word_length': 3.418974661166765, 'pattern_counts': {'":': 0, '<': 40, '<?xml version=': 0, '>': 42, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2403205', 'n_tokens_mistral': 2233, 'n_tokens_neox': 2175, 'n_words': 449}