package mclab.ide.callgraph;

import java.util.Arrays;
import java.util.List;

import ast.ASTNode;
import ast.AssignStmt;
import ast.ColonExpr;
import ast.Expr;
import ast.ExprStmt;
import ast.Function;
import ast.FunctionHandleExpr;
import ast.Name;
import ast.NameExpr;
import ast.ParameterizedExpr;
import ast.Program;
import ast.Script;
import ast.Stmt;
import ast.StringLiteralExpr;
import mclint.MatlabProgram;
import mclint.Project;
import mclint.util.AstUtil;
import natlab.LookupFile;
import natlab.toolkits.analysis.varorfun.VFAnalysis;
import natlab.toolkits.analysis.varorfun.VFPreorderAnalysis;
import natlab.toolkits.filehandling.FunctionOrScriptQuery;
import natlab.toolkits.path.FileEnvironment;
import natlab.utils.NodeFinder;
import nodecases.AbstractNodeCaseHandler;

public class CallgraphTransformer extends AbstractNodeCaseHandler {
  private static List<String> REFLECTIVE_FUNCTION_NAMES = Arrays.asList(
      "nargin", "nargout", "nargchk", "narginchk", "inputname");
  private static FunctionOrScriptQuery BASELINE_QUERY = LookupFile.getFunctionOrScriptQueryObject();
  private static FunctionOrScriptQuery getKindAnalysisEnvironment(Program node) {
    final FunctionOrScriptQuery env = new FileEnvironment(node.getFile())
        .getFunctionOrScriptQuery(node.getFile());
    return new FunctionOrScriptQuery() {
      @Override public boolean isFunctionOrScript(String name) {
        return env.isFunctionOrScript(name) || BASELINE_QUERY.isFunctionOrScript(name);
      }

      @Override public boolean isPackage(String name) {
        return env.isPackage(name) || BASELINE_QUERY.isPackage(name);
      }
    };
  }

  public static void instrument(Project project) {
    project.getMatlabPrograms().forEach(CallgraphTransformer::instrument);
  }

  public static void instrument(MatlabProgram program) {
    Program node = program.parse();
    VFAnalysis kindAnalysis = new VFPreorderAnalysis(node, getKindAnalysisEnvironment(node));
    kindAnalysis.analyze();
    node.analyze(new CallgraphTransformer(kindAnalysis, program.getPath().toString()));
  }

  private VFAnalysis kindAnalysis;
  private String relativePath;

  private CallgraphTransformer(VFAnalysis kindAnalysis, String relativePath) {
    this.kindAnalysis = kindAnalysis;
    this.relativePath = relativePath;
  }

  private String identifier(ASTNode node, String id) {
    return String.format("%s@%s:%d,%d", id, relativePath, node.getStartLine(), node.getStartColumn());
  }

  private String identifier(Name name) {
    return identifier(name, name.getID());
  }

  private String identifier(Function function) {
    return identifier(function, function.getName().getID());
  }

  private String identifier(Script script) {
    return identifier(script, script.getName());
  }

  private boolean isCall(ParameterizedExpr call) {
    return call.getTarget() instanceof NameExpr && isCall((NameExpr) call.getTarget());
  }

  private boolean isVar(ParameterizedExpr call) {
    return call.getTarget() instanceof NameExpr &&
        kindAnalysis.getResult(((NameExpr) call.getTarget()).getName()).isVariable() &&
        !(call.getParent() instanceof AssignStmt && ((AssignStmt) call.getParent()).getLHS() == call);
  }

  private boolean isCall(NameExpr call) {
    return kindAnalysis.getResult(call.getName()).isFunction() &&
        !REFLECTIVE_FUNCTION_NAMES.contains(call.getName().getID());
  }

  private static <T extends ASTNode> ast.List<T> concat(ast.List<T> list, ast.List<T> other) {
    other.forEach(list::addChild);
    return list;
  }

  private Expr wrapWithTraceCall(ParameterizedExpr call, boolean isVar) {
    Name target = ((NameExpr) call.getTarget()).getName();
    return new ParameterizedExpr(
        new NameExpr(new Name("mclab_callgraph_log_then_run")),
        concat(
          new ast.List<Expr>()
            .add(new StringLiteralExpr(isVar ? target.getID() : ""))
            .add(new StringLiteralExpr("call " + identifier(target) + "\\n"))
            .add(isVar ?
                new NameExpr(new Name(target.getID())) :
                new FunctionHandleExpr(new Name(target.getID()))),
          call.getArgs()
        )
    );
  }

  private Stmt entryPointLogStmt(String identifier) {
    return new ExprStmt(
        new ParameterizedExpr(
            new NameExpr(new Name("mclab_callgraph_log")),
            new ast.List<Expr>()
                .add(new StringLiteralExpr("enter " + identifier + "\\n"))
        )
    );
  }

  @Override public void caseASTNode(ASTNode node) {
    for (int i = 0; i < node.getNumChild(); ++i) {
      node.getChild(i).analyze(this);
    }
  }

  @Override public void caseFunction(Function f) {
    caseASTNode(f);
    f.getStmts().insertChild(entryPointLogStmt(identifier(f)), 0);
  }

  @Override public void caseScript(Script s) {
    caseASTNode(s);
    s.getStmts().insertChild(entryPointLogStmt(identifier(s)), 0);
  }

  @Override public void caseParameterizedExpr(ParameterizedExpr e) {
    e.getArgs().analyze(this);
    if (isCall(e)) {
      AstUtil.replace(e, wrapWithTraceCall(e, false /* isVar */));
    } else if (isVar(e)) {
      // Replace any colons with colon string literals; passing literal
      // colons to functions seems to confuse MATLAB.
      NodeFinder.find(ColonExpr.class, e.getArgs()).forEach(
          node -> AstUtil.replace(node, new StringLiteralExpr(":"))
      );
      AstUtil.replace(e, wrapWithTraceCall(e, true /* isVar */));
    }
  }

  @Override public void caseNameExpr(NameExpr e) {
    if (isCall(e)) {
      AstUtil.replace(e, wrapWithTraceCall(
          new ParameterizedExpr(e.fullCopy(), new ast.List<>()), false /* isVar */));
    }
  }
}

@@ -223,24 +223,47 @@
 HandleExpr;%0A
+import ast.LambdaExpr;%0A
 import ast.N
@@ -4919,32 +4919,504 @@
 r(s)), 0);%0A  %7D%0A%0A
+  @Override public void caseLambdaExpr(LambdaExpr e) %7B%0A    caseASTNode(e);%0A    AstUtil.replace(e.getBody(),%0A        new ParameterizedExpr(%0A            new NameExpr(new Name(%22mclab_callgraph_log_then_run%22)),%0A            new ast.List%3CExpr%3E()%0A                .add(new StringLiteralExpr(%22%22))%0A                .add(new StringLiteralExpr(%22enter %22 + identifier(e, %22%3Clambda%3E%22) + %22%5C%5Cn%22))%0A                .add(new LambdaExpr(new ast.List%3C%3E(), (Expr) e.getBody().fullCopy()))));%0A  %7D%0A%0A
   @Override publ

Add support for jump to function handles.