java-commons/muscript/src/main/java/io/gitlab/jfronny/muscript/ast/dynamic/ExprGroup.java

package io.gitlab.jfronny.muscript.ast.dynamic;

import io.gitlab.jfronny.muscript.ast.DynamicExpr;
import io.gitlab.jfronny.muscript.ast.Expr;
import io.gitlab.jfronny.muscript.ast.NullLiteral;
import io.gitlab.jfronny.muscript.ast.string.Concatenate;
import io.gitlab.jfronny.muscript.compiler.CodeLocation;
import io.gitlab.jfronny.muscript.compiler.ExprWriter;
import io.gitlab.jfronny.muscript.compiler.Order;
import io.gitlab.jfronny.muscript.data.Scope;
import io.gitlab.jfronny.muscript.data.dynamic.Dynamic;
import io.gitlab.jfronny.muscript.error.LocationalException;
import org.jetbrains.annotations.Nullable;

import java.io.IOException;
import java.util.LinkedList;
import java.util.List;
import java.util.stream.IntStream;
import java.util.stream.Stream;

public class ExprGroup extends DynamicExpr {
    private final List<Expr<?>> steps;
    private final DynamicExpr fin;
    private final @Nullable PackedArgs packedArgs;
    private final boolean fork;

    public ExprGroup(CodeLocation location, List<Expr<?>> expressions) {
        this(location, expressions, null);
    }

    public ExprGroup(CodeLocation location, List<Expr<?>> expressions, @Nullable PackedArgs packedArgs) {
        this(location, expressions, packedArgs, true);
    }

    public ExprGroup(CodeLocation location, List<Expr<?>> expressions, @Nullable PackedArgs packedArgs, boolean fork) {
        super(Order.Primary, location);
        this.steps = expressions.subList(0, expressions.size() - 1);
        this.fin = expressions.getLast().asDynamicExpr();
        this.packedArgs = packedArgs;
        this.fork = fork;
        if (!fork && packedArgs != null) throw new UnsupportedOperationException("packedArgs may only be used for forking ExprGroups");
    }

    public static Expr<?> of(CodeLocation location, List<Expr<?>> expressions) {
        return of(location, expressions, true);
    }

    public static Expr<?> of(CodeLocation location, List<Expr<?>> expressions, boolean fork) {
        if (expressions.isEmpty()) return new NullLiteral(location);
        if (!fork && expressions.size() == 1) return expressions.getFirst();
        return new ExprGroup(location, expressions, null, fork);
    }

    @Override
    public Dynamic get(Scope dataRoot) {
        Scope fork = this.fork ? dataRoot.fork() : dataRoot;
        if (this.fork && packedArgs != null) {
            List<Dynamic> from = new LinkedList<>();
            for (Call.Arg arg : packedArgs.from) {
                Dynamic data = arg.expr().get(dataRoot);
                if (arg.variadic()) from.addAll(data.asList().getValue());
                else from.add(data);
            }
            List<String> to = packedArgs.to;
            // Compare with Closure.get()
            int ac = from.size();
            int ae = to.size();
            if (packedArgs.variadic) ae--;
            if (ac < ae) throw new LocationalException(location, "Invoked with too few arguments (expected " + ae + " but got " + ac + ")");
            if (!packedArgs.variadic && ac > ae) throw new LocationalException(location, "Invoked with too many arguments (expected " + ae + " but got " + ac + ")");
            for (int i = 0; i < ae; i++) fork.set(to.get(i), from.get(i));
            if (packedArgs.variadic) {
                fork.set(to.get(to.size() - 1), IntStream.range(ae, ac)
                        .mapToObj(from::get)
                        .toList());
            }
        }
        for (Expr<?> step : steps) {
            step.get(fork);
        }
        return fin.get(fork);
    }

    @Override
    public DynamicExpr optimize() {
        List<Expr<?>> exprs = Stream.concat(
                steps.stream()
                        .map(Expr::optimize)
                        .flatMap(Expr::extractSideEffects)
                        .map(Expr::optimize),
                Stream.of(fin.optimize())
        ).toList();
        return packedArgs == null
                ? of(location, exprs, fork).asDynamicExpr()
                : new ExprGroup(location, exprs, packedArgs, fork);
    }

    @Override
    public Stream<Expr<?>> extractSideEffects() {
        return fork ? Stream.of(this) : Stream.of(of(location, stream().flatMap(Expr::extractSideEffects).toList()));
    }

    @Override
    public void decompile(ExprWriter writer) throws IOException {
        if (!fork) {
            // Use string concatenation since the result is ignored either way
            stream()
                    .reduce((left, right) -> new Concatenate(location, left.asStringExpr(), right.asStringExpr()))
                    .orElseGet(() -> new NullLiteral(location))
                    .optimize()
                    .decompile(writer);
        } else {
            new Call(
                    location,
                    new Closure(
                            location,
                            packedArgs == null
                                    ? List.of()
                                    : packedArgs.to,
                            stream().toList(),
                            false
                    ),
                    packedArgs == null ? List.of() : packedArgs.from
            ).decompile(writer);
        }
    }

    public Stream<Expr<?>> stream() {
        return Stream.concat(steps.stream(), Stream.of(fin));
    }

    public record PackedArgs(List<Call.Arg> from, List<String> to, boolean variadic) {}
}