<commit_before>package yokohama.unit.ast;

import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Optional;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import lombok.RequiredArgsConstructor;
import yokohama.unit.util.Lists;
import yokohama.unit.util.Pair;

@RequiredArgsConstructor
public class ChoiceCollectVisitor extends 
        StreamVisitorTemplate<Pair<List<Ident>, List<List<Expr>>>> {
    final List<Table> tables;

    @Override
    public Stream<Pair<List<Ident>, List<List<Expr>>>> visitChoiceBinding(ChoiceBinding choiceBinding) {
        List<Ident> idents = Arrays.asList(choiceBinding.getName());
        List<List<Expr>> ess = Lists.map(
                choiceBinding.getChoices(),
                choice -> Arrays.asList(choice));
        return Stream.of(new Pair<>(idents, ess));
    }

    @Override
    public Stream<Pair<List<Ident>, List<List<Expr>>>> visitTableBinding(
            TableBinding tableBinding) {
        List<Ident> idents = tableBinding.getIdents();
        String name = tableBinding.getName();
        List<List<Expr>> ess;
        switch (tableBinding.getType()) {
            case INLINE:
                Optional<Table> tableOpt =
                        tables.stream()
                                .filter(table -> table.getName().equals(name))
                                .findFirst();
                if (tableOpt.isPresent()) {
                    Table table = tableOpt.get();
                    List<Ident> header = table.getHeader();
                    ess = table.getRows().stream()
                            .map(row ->
                                    Pair.zip(header, row.getCells())
                                            .stream()
                                            .collect(
                                                    () -> new HashMap<Ident, Cell>(),
                                                    (m, p) -> m.put(
                                                            p.getFirst(),
                                                            p.getSecond()),
                                                    (m1, m2) -> m1.putAll(m2)))
                            .map(m -> {
                                return header.stream()
                                        .map(ident -> m.get(ident))
                                        .map(cell -> {
                                            return cell.accept(
                                                    expr -> expr.getExpr(),
                                                    pred -> {
                                                        throw new AstException(
                                                                "invalid reference to predicate cell",
                                                                tableBinding.getSpan());<commit_msg>Fix: only columns mathcing idents should be selected<commit_after>package yokohama.unit.ast;

import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Optional;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import lombok.RequiredArgsConstructor;
import yokohama.unit.util.Lists;
import yokohama.unit.util.Pair;

@RequiredArgsConstructor
public class ChoiceCollectVisitor extends 
        StreamVisitorTemplate<Pair<List<Ident>, List<List<Expr>>>> {
    final List<Table> tables;

    @Override
    public Stream<Pair<List<Ident>, List<List<Expr>>>> visitChoiceBinding(ChoiceBinding choiceBinding) {
        List<Ident> idents = Arrays.asList(choiceBinding.getName());
        List<List<Expr>> ess = Lists.map(
                choiceBinding.getChoices(),
                choice -> Arrays.asList(choice));
        return Stream.of(new Pair<>(idents, ess));
    }

    @Override
    public Stream<Pair<List<Ident>, List<List<Expr>>>> visitTableBinding(
            TableBinding tableBinding) {
        List<Ident> idents = tableBinding.getIdents();
        String name = tableBinding.getName();
        List<List<Expr>> ess;
        switch (tableBinding.getType()) {
            case INLINE:
                Optional<Table> tableOpt =
                        tables.stream()
                                .filter(table -> table.getName().equals(name))
                                .findFirst();
                if (tableOpt.isPresent()) {
                    Table table = tableOpt.get();
                    List<Ident> header = table.getHeader();
                    ess = table.getRows().stream()
                            .map(row ->
                                    Pair.zip(header, row.getCells())
                                            .stream()
                                            .collect(
                                                    () -> new HashMap<Ident, Cell>(),
                                                    (m, p) -> m.put(
                                                            p.getFirst(),
                                                            p.getSecond()),
                                                    (m1, m2) -> m1.putAll(m2)))
                            .map(m -> {
                                return idents.stream()
                                        .map(ident -> m.get(ident))
                                        .map(cell -> {
                                            return cell.accept(
                                                    expr -> expr.getExpr(),
                                                    pred -> {
                                                        throw new AstException(
                                                                "invalid reference to predicate cell",
                                                                tableBinding.getSpan());
{'fraction_non_alphanumeric': 0.09670816987890159, 'fraction_numerical': 0.0013644891693672181, 'mean_word_length': 1.0206753962784287, 'pattern_counts': {'":': 0, '<': 55, '<?xml version=': 0, '>': 77, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 58, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3629973', 'n_tokens_mistral': 1214, 'n_tokens_neox': 1175, 'n_words': 246}