ModulesMapper.java

package scotch.compiler.parser;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import org.antlr.v4.runtime.ParserRuleContext;
import scotch.compiler.parser.ScotchParser.ModuleContext;
import scotch.compiler.parser.ScotchParser.ModulesContext;

public class ModulesMapper {

    public Map<String, List<ModuleContext>> map(List<? extends ParserRuleContext> ctx) {
        return new State().map(ctx);
    }

    private final class State extends ScotchParserBaseVisitor<Void> {

        private final Map<String, List<ModuleContext>> moduleMap;

        public State() {
            moduleMap = new HashMap<>();
        }

        public Map<String, List<ModuleContext>> map(List<? extends ParserRuleContext> ctxs) {
            ctxs.forEach(ctx -> ctx.accept(this));
            return moduleMap;
        }

        @Override
        public Void visitModule(ModuleContext ctx) {
            moduleMap.computeIfAbsent(ctx.moduleName().getText(), k -> new ArrayList<>()).add(ctx);
            return null;
        }

        @Override
        public Void visitModules(ModulesContext ctx) {
            ctx.module().forEach(module -> module.accept(this));
            return null;
        }

        @Override
        protected Void defaultResult() {
            return null;
        }
    }
}