package io.github.vmzakharov.ecdataframe.grammar;
import io.github.vmzakharov.ecdataframe.dsl.*;
import io.github.vmzakharov.ecdataframe.dsl.value.DoubleValue;
import io.github.vmzakharov.ecdataframe.dsl.value.LongValue;
import io.github.vmzakharov.ecdataframe.dsl.value.StringValue;
import org.antlr.v4.runtime.Token;
import org.eclipse.collections.api.factory.Stacks;
import org.eclipse.collections.api.list.ListIterable;
import org.eclipse.collections.api.list.MutableList;
import org.eclipse.collections.api.stack.MutableStack;
import org.eclipse.collections.impl.factory.Lists;
import org.eclipse.collections.impl.utility.ListIterate;
public class ModelScriptTreeBuilderVisitor
extends ModelScriptBaseVisitor<Expression>
{
private final MutableStack<Script> scriptStack = Stacks.mutable.of();
public ModelScriptTreeBuilderVisitor()
{
this.scriptStack.push(new AnonymousScript());
}
@Override
public Expression visitFreeExp(ModelScriptParser.FreeExpContext ctx)
{
Expression freeExpression = this.visit(ctx.expr());
return this.addStatementToCurrentScriptContext(freeExpression);
}
@Override
public Expression visitAssignExpr(ModelScriptParser.AssignExprContext ctx)
{
String varName = ctx.ID().getText();
if (ctx.expr() == null)
{
throw new RuntimeException("Malformed assignment statement");
}
Expression expression = this.visit(ctx.expr());
return this.addStatementToCurrentScriptContext(new AssingExpr(varName, expression));
}
private Expression addStatementToCurrentScriptContext(Expression expression)
{
return this.getScript().addStatement(expression);
}
@Override
public Expression visitPropertyPathExpr(ModelScriptParser.PropertyPathExprContext ctx)
{
return new PropertyPathExpr(ListIterate.collect(ctx.ID(), Object::toString));
}
@Override
public Expression visitAliasExpr(ModelScriptParser.AliasExprContext ctx)
{
return new AliasExpr(ctx.ID().toString(), this.visit(ctx.expr()));
}
@Override
public Expression visitStringLiteralExpr(ModelScriptParser.StringLiteralExprContext ctx)
{
return new ConstExpr(new StringValue(this.stripQuotes(ctx.STRING().getText())));
}
@Override
public Expression visitVarExpr(ModelScriptParser.VarExprContext ctx)
{
return new VarExpr(ctx.ID().getText());
}
@Override
public Expression visitIntLiteralExpr(ModelScriptParser.IntLiteralExprContext ctx)
{
return new ConstExpr(new LongValue(Integer.parseInt(ctx.INT().getText())));
}
@Override
public Expression visitDoubleLiteralExpr(ModelScriptParser.DoubleLiteralExprContext ctx)
{
return new ConstExpr(new DoubleValue(Double.parseDouble(ctx.DOUBLE().getText())));
}
@Override
public Expression visitAddSubExpr(ModelScriptParser.AddSubExprContext ctx)
{
return this.visitBinaryOperation(ctx.expr(0), ctx.expr(1), ctx.op);
}
@Override
public Expression visitMulDivExpr(ModelScriptParser.MulDivExprContext ctx)
{
return this.visitBinaryOperation(ctx.expr(0), ctx.expr(1), ctx.op);
}
private Expression visitBinaryOperation(ModelScriptParser.ExprContext exprContext1, ModelScriptParser.ExprContext exprContext2, Token opToken)
{
ArithmeticOp operation = null;
switch (opToken.getType())
{
case ModelScriptParser.MUL: operation = ArithmeticOp.MULTIPLY; break;
case ModelScriptParser.DIV: operation = ArithmeticOp.DIVIDE; break;
case ModelScriptParser.ADD: operation = ArithmeticOp.ADD; break;
case ModelScriptParser.SUB: operation = ArithmeticOp.SUBTRACT; break;
}
return new BinaryExpr(this.visit(exprContext1), this.visit(exprContext2), operation);
}
@Override
public Expression visitAndExpr(ModelScriptParser.AndExprContext ctx)
{
return this.visitBooleanOperation(ctx.expr(0), ctx.expr(1), ctx.op);
}
@Override
public Expression visitOrExpr(ModelScriptParser.OrExprContext ctx)
{
return this.visitBooleanOperation(ctx.expr(0), ctx.expr(1), ctx.op);
}
private Expression visitBooleanOperation(ModelScriptParser.ExprContext exprContext1, ModelScriptParser.ExprContext exprContext2, Token opToken)
{
BooleanOp operation = null;
switch (opToken.getType())
{
case ModelScriptParser.AND:
operation = BooleanOp.AND;
break;
case ModelScriptParser.OR:
operation = BooleanOp.OR;
break;
case ModelScriptParser.XOR:
operation = BooleanOp.XOR;
break;
}
return new BinaryExpr(this.visit(exprContext1), this.visit(exprContext2), operation);
}
@Override
public Expression visitCompareExpr(ModelScriptParser.CompareExprContext ctx)
{
ComparisonOp operation = null;
switch (ctx.op.getType())
{
case ModelScriptParser.GT : operation = ComparisonOp.GT ; break;
case ModelScriptParser.GTE: operation = ComparisonOp.GTE; break;
case ModelScriptParser.LT : operation = ComparisonOp.LT ; break;
case ModelScriptParser.LTE: operation = ComparisonOp.LTE; break;
case ModelScriptParser.EQ : operation = ComparisonOp.EQ ; break;
case ModelScriptParser.NE : operation = ComparisonOp.NE ; break;
}
return new BinaryExpr(this.visit(ctx.expr(0)), this.visit(ctx.expr(1)), operation);
}
@Override
public Expression visitParenExpr(ModelScriptParser.ParenExprContext ctx)
{
return super.visit(ctx.expr());
}
@Override
public Expression visitFunctionDeclarationExpr(ModelScriptParser.FunctionDeclarationExprContext ctx)
{
String functionName = ctx.ID().toString();
MutableList<String> parameterNames;
if (ctx.idList() == null)
{
parameterNames = Lists.mutable.of();
}
else
{
parameterNames = ListIterate.collect(ctx.idList().ID(), Object::toString);
}
FunctionScript functionScript = new FunctionScript(functionName, parameterNames);
this.pushScript(functionScript);
ListIterate.forEach(ctx.statementSequence().statement(), this::visit);
this.popScript();
this.getAsAnonymousScript().addFunctionScript(functionScript);
return null;
}
@Override
public Expression visitUnaryMinusExpr(ModelScriptParser.UnaryMinusExprContext ctx)
{
return new UnaryExpr(UnaryOp.MINUS, this.visit(ctx.expr()));
}
@Override
public Expression visitNotExpr(ModelScriptParser.NotExprContext ctx)
{
return new UnaryExpr(UnaryOp.NOT, this.visit(ctx.expr()));
}
@Override
public Expression visitInExpr(ModelScriptParser.InExprContext ctx)
{
return new BinaryExpr(
this.visit(ctx.expr(0)),
this.visit(ctx.expr(1)),
ctx.IN() == null ? ContainsOp.NOT_IN : ContainsOp.IN);
}
@Override
public Expression visitIsEmptyExpr(ModelScriptParser.IsEmptyExprContext ctx)
{
return new UnaryExpr(UnaryOp.IS_EMPTY, this.visit(ctx.expr()));
}
@Override
public Expression visitIsNotEmptyExpr(ModelScriptParser.IsNotEmptyExprContext ctx)
{
return new UnaryExpr(UnaryOp.IS_NOT_EMPTY, this.visit(ctx.expr()));
}
@Override
public Expression visitVectorExpr(ModelScriptParser.VectorExprContext ctx)
{
ListIterable<Expression> elements;
if (ctx.exprList() == null)
{
elements = Lists.immutable.of();
}
else
{
elements = ListIterate.collect(ctx.exprList().expr(), this::visit);
}
return new VectorExpr(elements);
}
@Override
public Expression visitIndexVectorExpr(ModelScriptParser.IndexVectorExprContext ctx)
{
return new IndexExpr(
this.visit(ctx.expr(0)),
this.visit(ctx.expr(1))
);
}
@Override
public Expression visitFunctionCallExpr(ModelScriptParser.FunctionCallExprContext ctx)
{
String functionName = ctx.ID().toString();
ListIterable<Expression> parameters;
if (ctx.exprList() == null)
{
parameters = Lists.immutable.of();
}
else
{
parameters = ListIterate.collect(ctx.exprList().expr(), this::visit);
}
return new FunctionCallExpr(functionName, parameters);
}
@Override
public Expression visitTernaryExpr(ModelScriptParser.TernaryExprContext ctx)
{
return new IfElseExpr(this.visit(ctx.condExpr), this.visit(ctx.ifExpr), this.visit(ctx.elseExpr), true);
}
@Override
public Expression visitConditionExpr(ModelScriptParser.ConditionExprContext ctx)
{
Expression condition = this.visit(ctx.expr());
Script ifScript = this.processStatementSequence(ctx.ifBody);
if (ctx.elseBody == null)
{
this.addStatementToCurrentScriptContext(new IfElseExpr(condition, ifScript));
}
else
{
Script elseScript = this.processStatementSequence(ctx.elseBody);
this.addStatementToCurrentScriptContext(new IfElseExpr(condition, ifScript, elseScript, false));
}
return null;
}
private Script processStatementSequence(ModelScriptParser.StatementSequenceContext ctx)
{
this.pushScript(new StatementSequenceScript());
ListIterate.forEach(ctx.statement(), this::visit);
return this.popScript();
}
@Override
public Expression visitProjectionStatement(ModelScriptParser.ProjectionStatementContext ctx)
{
ListIterable<Expression> projectionList = ListIterate.collect(ctx.exprList().expr(), this::visit);
ProjectionExpr projectionExpr = (ctx.expr() == null) ?
new ProjectionExpr(projectionList) :
new ProjectionExpr(projectionList, this.visit(ctx.expr()));
return this.addStatementToCurrentScriptContext(projectionExpr);
}
private String stripQuotes(String aString)
{
if (aString.length() < 2) return aString;
if (aString.charAt(0) == '"' && aString.charAt(aString.length()-1) == '"')
{
return aString.substring(1, aString.length()-1);
}
if (aString.charAt(0) == '\'' && aString.charAt(aString.length()-1) == '\'')
{
return aString.substring(1, aString.length()-1);
}
return aString;
}
public Script getScript()
{
return this.scriptStack.peek();
}
public AnonymousScript getAsAnonymousScript()
{
return (AnonymousScript) this.scriptStack.peek();
}
private Script pushScript(Script newScript)
{
this.scriptStack.push(newScript);
return newScript;
}
private Script popScript()
{
return this.scriptStack.pop();
}
}