/*
 * Copyright (c) 2002 - 2012 IBM Corporation.
 * All rights reserved. This program and the accompanying materials
 * are made available under the terms of the Eclipse Public License v1.0
 * which accompanies this distribution, and is available at
 * http://www.eclipse.org/legal/epl-v10.html
 *
 * Contributors:
 *     IBM Corporation - initial API and implementation
 */
package com.ibm.wala.cast.js.callgraph.fieldbased.flowgraph;

import com.ibm.wala.cast.ir.ssa.AstGlobalRead;
import com.ibm.wala.cast.ir.ssa.AstGlobalWrite;
import com.ibm.wala.cast.ir.ssa.AstLexicalAccess.Access;
import com.ibm.wala.cast.ir.ssa.AstLexicalRead;
import com.ibm.wala.cast.ir.ssa.AstLexicalWrite;
import com.ibm.wala.cast.ir.ssa.AstPropertyRead;
import com.ibm.wala.cast.ir.ssa.AstPropertyWrite;
import com.ibm.wala.cast.js.callgraph.fieldbased.JSMethodInstructionVisitor;
import com.ibm.wala.cast.js.callgraph.fieldbased.flowgraph.vertices.CreationSiteVertex;
import com.ibm.wala.cast.js.callgraph.fieldbased.flowgraph.vertices.FuncVertex;
import com.ibm.wala.cast.js.callgraph.fieldbased.flowgraph.vertices.VarVertex;
import com.ibm.wala.cast.js.callgraph.fieldbased.flowgraph.vertices.Vertex;
import com.ibm.wala.cast.js.callgraph.fieldbased.flowgraph.vertices.VertexFactory;
import com.ibm.wala.cast.js.ipa.callgraph.JSCallGraphUtil;
import com.ibm.wala.cast.js.ipa.callgraph.JSSSAPropagationCallGraphBuilder;
import com.ibm.wala.cast.js.ssa.JavaScriptInvoke;
import com.ibm.wala.cast.js.ssa.PrototypeLookup;
import com.ibm.wala.cast.js.ssa.SetPrototype;
import com.ibm.wala.cast.js.types.JavaScriptMethods;
import com.ibm.wala.cast.js.types.JavaScriptTypes;
import com.ibm.wala.cast.js.util.Util;
import com.ibm.wala.cast.loader.AstMethod;
import com.ibm.wala.cast.loader.AstMethod.LexicalInformation;
import com.ibm.wala.cast.types.AstMethodReference;
import com.ibm.wala.classLoader.IClass;
import com.ibm.wala.classLoader.IMethod;
import com.ibm.wala.ipa.callgraph.IAnalysisCacheView;
import com.ibm.wala.ipa.cha.IClassHierarchy;
import com.ibm.wala.ssa.IR;
import com.ibm.wala.ssa.SSAGetCaughtExceptionInstruction;
import com.ibm.wala.ssa.SSAGetInstruction;
import com.ibm.wala.ssa.SSAInstruction;
import com.ibm.wala.ssa.SSANewInstruction;
import com.ibm.wala.ssa.SSAPhiInstruction;
import com.ibm.wala.ssa.SSAPutInstruction;
import com.ibm.wala.ssa.SSAReturnInstruction;
import com.ibm.wala.ssa.SSAThrowInstruction;
import com.ibm.wala.types.TypeReference;
import com.ibm.wala.util.collections.Iterator2Iterable;
import com.ibm.wala.util.intset.EmptyIntSet;
import com.ibm.wala.util.intset.IntSet;

/**
 * Class for building intra-procedural flow graphs for a given class hierarchy.
 *
 * @author mschaefer
 */
public class FlowGraphBuilder {
  private final IClassHierarchy cha;
  private final IAnalysisCacheView cache;
  private final boolean supportFullPointerAnalysis;

  public FlowGraphBuilder(
      IClassHierarchy cha, IAnalysisCacheView cache, boolean supportPointerAnalysis) {
    this.cha = cha;
    this.cache = cache;
    this.supportFullPointerAnalysis = supportPointerAnalysis;
  }

  /**
   * This is the main entry point of the flow graph builder.
   *
   * <p>It creates a new, empty flow graph, adds nodes for a small number of special primitive
   * functions such as {@code Object} and {@code Function} and sets up flow edges to make them flow
   * into the corresponding global variables. Then it iterates over all functions in the class
   * hierarchy and all their IR instructions, and adds the flow edges induced by these instructions.
   *
   * @return the completed flow graph
   */
  public FlowGraph buildFlowGraph() {
    FlowGraph flowgraph = new FlowGraph();

    addPrimitives(flowgraph);

    visitProgram(flowgraph);

    return flowgraph;
  }

  public void visitProgram(FlowGraph flowgraph) {
    for (IClass klass : cha) {
      for (IMethod method : klass.getDeclaredMethods()) {
        if (method.getDescriptor().equals(AstMethodReference.fnDesc)) {
          visitFunction(flowgraph, method);
        }
      }
    }
  }

  public void visitFunction(FlowGraph flowgraph, IMethod method) {
    {
      IR ir = cache.getIR(method);
      FlowGraphSSAVisitor visitor = new FlowGraphSSAVisitor(ir, flowgraph);

      // first visit normal instructions
      SSAInstruction[] normalInstructions = ir.getInstructions();
      for (int i = 0; i < normalInstructions.length; ++i)
        if (normalInstructions[i] != null) {
          visitor.instructionIndex = i;
          normalInstructions[i].visit(visitor);
        }

      // now visit phis and catches
      visitor.instructionIndex = -1;
      for (SSAInstruction inst : Iterator2Iterable.make(ir.iteratePhis())) inst.visit(visitor);

      for (SSAInstruction inst : Iterator2Iterable.make(ir.iterateCatchInstructions()))
        inst.visit(visitor);
    }
  }

  // primitive functions that are treated specially
  private static String[] primitiveFunctions = {
    "Object", "Function", "Array", "StringObject", "NumberObject", "BooleanObject", "RegExp"
  };

  /**
   * Add flows from the special primitive functions to the corresponding global variables.
   *
   * @param flowgraph the flow graph under construction
   */
  private void addPrimitives(FlowGraph flowgraph) {
    VertexFactory factory = flowgraph.getVertexFactory();
    for (String pf : primitiveFunctions) {
      TypeReference typeref = TypeReference.findOrCreate(JavaScriptTypes.jsLoader, 'L' + pf);
      IClass klass = cha.lookupClass(typeref);
      String prop = pf.endsWith("Object") ? pf.substring(0, pf.length() - 6) : pf;
      flowgraph.addEdge(factory.makeFuncVertex(klass), factory.makePropVertex(prop));
    }
  }

  /**
   * Visitor class that does the heavy lifting (such as it is) of flow graph construction, adding
   * flow graph edges for every instruction in the method IR.
   *
   * <p>The only slightly tricky thing are assignments to exposed variables inside their defining
   * function. In the IR, they initially appear as normal SSA variable assignments, without any
   * indication of their lexical nature. The normal call graph construction logic does something
   * convoluted to fix this up later when an actual lexical access is encountered.
   *
   * <p>We use a much simpler approach. Whenever we see an assignment <code>v<sub>i</sub> = e</code>
   * , we ask the enclosing function whether <code>v<sub>i</sub></code> is an exposed variable. If
   * it is, we determine its source-level names <code>
   * x<sub>1</sub>, x<sub>2</sub>, ..., x<sub>n</sub></code>, and then add edges corresponding to
   * lexical writes of <code>v<sub>i</sub></code> into all the <code>x<sub>j</sub></code>.
   *
   * @author mschaefer
   */
  private class FlowGraphSSAVisitor extends JSMethodInstructionVisitor {
    // index of the instruction currently visited; -1 if the instruction isn't a normal instruction
    public int instructionIndex = -1;

    // flow graph being built
    private final FlowGraph flowgraph;

    // vertex factory to use for constructing new vertices
    private final VertexFactory factory;

    // lexical information about the current function
    private final LexicalInformation lexicalInfo;

    // the set of SSA variables in the current function that are accessed by nested functions
    private final IntSet exposedVars;

    // the IR of the current function
    private final IR ir;

    // the function vertex corresponding to the current function
    private final FuncVertex func;

    public FlowGraphSSAVisitor(IR ir, FlowGraph flowgraph) {
      super(ir.getMethod(), ir.getSymbolTable(), cache.getDefUse(ir));
      this.ir = ir;
      this.flowgraph = flowgraph;
      this.factory = flowgraph.getVertexFactory();
      this.func = factory.makeFuncVertex(ir.getMethod().getDeclaringClass());
      if (method instanceof AstMethod) {
        this.lexicalInfo = ((AstMethod) method).lexicalInfo();
        this.exposedVars = lexicalInfo.getAllExposedUses();
      } else {
        this.lexicalInfo = null;
        this.exposedVars = EmptyIntSet.instance;
      }
    }

    // add extra flow from v_def to every lexical variable it may correspond to at source-level
    private void handleLexicalDef(int def) {
      assert def != -1;
      if (instructionIndex != -1 && exposedVars.contains(def)) {
        VarVertex v = factory.makeVarVertex(func, def);
        for (String localName : ir.getLocalNames(instructionIndex, def))
          flowgraph.addEdge(
              v, factory.makeLexicalAccessVertex(lexicalInfo.getScopingName(), localName));
      }
    }

    @Override
    public void visitPhi(SSAPhiInstruction phi) {
      int n = phi.getNumberOfUses();
      VarVertex w = factory.makeVarVertex(func, phi.getDef());
      for (int i = 0; i < n; ++i) {
        VarVertex v = factory.makeVarVertex(func, phi.getUse(i));
        flowgraph.addEdge(v, w);
      }
    }

    @Override
    public void visitPrototypeLookup(PrototypeLookup proto) {
      // treat it simply as an assignment
      flowgraph.addEdge(
          factory.makeVarVertex(func, proto.getUse(0)),
          factory.makeVarVertex(func, proto.getDef()));
      handleLexicalDef(proto.getDef());
    }

    private void visitPut(int val, String propName) {
      Vertex v = factory.makeVarVertex(func, val), w = factory.makePropVertex(propName);
      flowgraph.addEdge(v, w);
    }

    @Override
    public void visitPut(SSAPutInstruction put) {
      visitPut(put.getVal(), put.getDeclaredField().getName().toString());
    }

    @Override
    public void visitSetPrototype(SetPrototype instruction) {
      visitPut(instruction.getUse(1), "prototype");
    }

    @Override
    public void visitAstGlobalWrite(AstGlobalWrite instruction) {
      String propName = instruction.getDeclaredField().getName().toString();

      // hack to account for global variables
      assert propName.startsWith("global ");
      propName = propName.substring("global ".length());

      visitPut(instruction.getVal(), propName);
    }

    @Override
    public void visitPropertyWrite(AstPropertyWrite pw) {
      int p = pw.getMemberRef();
      if (symtab.isConstant(p)) {
        String pn = JSCallGraphUtil.simulateToStringForPropertyNames(symtab.getConstantValue(p));

        Vertex v = factory.makeVarVertex(func, pw.getValue()), w = factory.makePropVertex(pn);
        flowgraph.addEdge(v, w);
      }
    }

    @Override
    public void visitAstLexicalWrite(AstLexicalWrite lw) {
      for (Access acc : lw.getAccesses()) {
        Vertex v = factory.makeVarVertex(func, acc.valueNumber),
            w = factory.makeLexicalAccessVertex(acc.variableDefiner, acc.variableName);
        flowgraph.addEdge(v, w);
      }
    }

    @Override
    public void visitGet(SSAGetInstruction get) {
      String propName = get.getDeclaredField().getName().toString();
      if (propName.startsWith("global ")) propName = propName.substring("global ".length());

      Vertex v = factory.makePropVertex(propName), w = factory.makeVarVertex(func, get.getDef());
      flowgraph.addEdge(v, w);
      handleLexicalDef(get.getDef());
    }

    @Override
    public void visitAstGlobalRead(AstGlobalRead instruction) {
      if (supportFullPointerAnalysis
          && instruction
              .getGlobalName()
              .endsWith(JSSSAPropagationCallGraphBuilder.GLOBAL_OBJ_VAR_NAME)) {
        Vertex lval = factory.makeVarVertex(func, instruction.getDef());
        flowgraph.addEdge(factory.global(), lval);
      } else {
        visitGet(instruction);
      }
    }

    @Override
    public void visitPropertyRead(AstPropertyRead pr) {
      int p = pr.getMemberRef();
      if (symtab.isConstant(p)) {
        String pn = JSCallGraphUtil.simulateToStringForPropertyNames(symtab.getConstantValue(p));
        Vertex v = factory.makePropVertex(pn), w = factory.makeVarVertex(func, pr.getDef());
        flowgraph.addEdge(v, w);
      }

      IntSet argVns = Util.getArgumentsArrayVns(ir, du);
      if (argVns.contains(pr.getObjectRef())) {
        Vertex v = factory.makeArgVertex(func), w = factory.makeVarVertex(func, pr.getDef());
        flowgraph.addEdge(v, w);
      }

      handleLexicalDef(pr.getDef());
    }

    @Override
    public void visitAstLexicalRead(AstLexicalRead lr) {
      for (Access acc : lr.getAccesses()) {
        Vertex v = factory.makeLexicalAccessVertex(acc.variableDefiner, acc.variableName),
            w = factory.makeVarVertex(func, acc.valueNumber);
        flowgraph.addEdge(v, w);
        handleLexicalDef(acc.valueNumber);
      }
    }

    @Override
    public void visitReturn(SSAReturnInstruction ret) {
      Vertex v = factory.makeVarVertex(func, ret.getResult()), w = factory.makeRetVertex(func);
      flowgraph.addEdge(v, w);
    }

    @Override
    public void visitThrow(SSAThrowInstruction thr) {
      Vertex v = factory.makeVarVertex(func, thr.getException()), w = factory.makeUnknownVertex();
      flowgraph.addEdge(v, w);
    }

    @Override
    public void visitGetCaughtException(SSAGetCaughtExceptionInstruction katch) {
      Vertex v = factory.makeUnknownVertex(), w = factory.makeVarVertex(func, katch.getDef());
      flowgraph.addEdge(v, w);
    }

    @Override
    public void visitJavaScriptInvoke(JavaScriptInvoke invk) {
      flowgraph.addEdge(
          factory.makeUnknownVertex(), factory.makeVarVertex(func, invk.getException()));

      // check whether this invoke corresponds to a function expression/declaration
      // flow callee variable into callee vertex
      if (invk.getDeclaredTarget().equals(JavaScriptMethods.ctorReference)) {

        flowgraph.addEdge(
            factory.makeVarVertex(func, invk.getFunction()), factory.makeCallVertex(func, invk));

        if (isFunctionConstructorInvoke(invk)) {
          // second parameter is function name
          String fn_name = symtab.getStringValue(invk.getUse(1));

          // find the function being defined here
          IClass klass =
              cha.lookupClass(TypeReference.findOrCreate(JavaScriptTypes.jsLoader, fn_name));
          if (klass == null) {
            System.err.println(
                "cannot find "
                    + fn_name
                    + " at "
                    + ((AstMethod) ir.getMethod())
                        .getSourcePosition(
                            ir.getCallInstructionIndices(invk.getCallSite()).intIterator().next()));
            return;
          }

          IMethod fn = klass.getMethod(AstMethodReference.fnSelector);
          FuncVertex fnVertex = factory.makeFuncVertex(klass);

          // function flows into its own v1 variable
          flowgraph.addEdge(fnVertex, factory.makeVarVertex(fnVertex, 1));

          // flow parameters into local variables
          for (int i = 1; i < fn.getNumberOfParameters(); ++i)
            flowgraph.addEdge(
                factory.makeParamVertex(fnVertex, i), factory.makeVarVertex(fnVertex, i + 1));

          // flow function into result variable
          flowgraph.addEdge(fnVertex, factory.makeVarVertex(func, invk.getDef()));

        } else if (supportFullPointerAnalysis) {

          CreationSiteVertex cs =
              factory.makeCreationSiteVertex(method, invk.iIndex(), JavaScriptTypes.Object);

          // flow creation site into result of new call
          flowgraph.addEdge(cs, factory.makeVarVertex(func, invk.getDef()));

          // also passed as 'this' to constructor
          if (invk.getNumberOfPositionalParameters() > 1) {
            flowgraph.addEdge(cs, factory.makeVarVertex(func, invk.getUse(0)));
          }
        }

      } else {
        // check whether it is a method call
        if (invk.getDeclaredTarget().equals(JavaScriptMethods.dispatchReference)) {
          // we only handle method calls with constant names
          if (symtab.isConstant(invk.getFunction())) {
            String pn =
                JSCallGraphUtil.simulateToStringForPropertyNames(
                    symtab.getConstantValue(invk.getFunction()));
            // flow callee property into callee vertex
            flowgraph.addEdge(factory.makePropVertex(pn), factory.makeCallVertex(func, invk));
          }
        } else {
          // this case is simpler: just flow callee variable into callee vertex
          flowgraph.addEdge(
              factory.makeVarVertex(func, invk.getFunction()), factory.makeCallVertex(func, invk));
        }
      }
      handleLexicalDef(invk.getDef());
    }

    @Override
    public void visitNew(SSANewInstruction invk) {
      if (supportFullPointerAnalysis) {
        // special case for supporting full pointer analysis
        // some core objects in the prologue (and the arguments array objects) get created with
        // 'new'
        CreationSiteVertex cs =
            factory.makeCreationSiteVertex(method, invk.iIndex(), invk.getConcreteType());

        // flow creation site into result of new call
        flowgraph.addEdge(cs, factory.makeVarVertex(func, invk.getDef()));
      }
    }
  }
}
