package tip.analysis

import tip.ast
import tip.ast._
import tip.lattices._
import tip.graph._
import tip.solvers._
import tip.graph.NodeOps._

object SignAnalysisCommons {
  val returnId = AIdentifier(s"#result", Loc(0, 0))()
}

/**
 * Transfer functions for sign analysis (intraprocedural only).
 */
trait IntraprocSignAnalysisTransferFunctions {

  type Sign = SignLattice.SignElement.Value

  import SignLattice.SignElement._

  val domain: Set[GNode[AstNode]]

  val declaredVars = domain.map {_.declaredIdsIncludingParams}.flatten

  val statelattice = new MapLattice(declaredVars, SignLattice)

  /**
   * The transfer functions.
   */
  def transfer(n: GNode[AstNode], s: statelattice.Element): statelattice.Element = {
    NoPointer.checkCfgNodeLanguageRestriction(n)
    import NoPointer._
    n match {
      case r: GRealNode[AstNode] =>
        r.data match {

          // var declarations
          case varr: AVarStmt =>
            varr.declIds.foldLeft(s)( (acc, id) => acc + (id -> Top) )

          // assignments
          case ass: AAssignStmt =>
            val id = ass.leftId
            val vdef = id.meta.definition match {
              case Some(x: AIdentifier) => x
              case Some(y) => unsupportedError(y)
              case _ => throw new IllegalArgumentException(s"No definition found for $id")
            }
            try {
              s + (vdef -> statelattice.l.eval(ass.right, s))
            }
            catch {
              case e: Exception => unsupportedError(ass)
            }

          // all others: like no-ops
          case _ => s
        }
      case _ => s
    }
  }

  private def unsupportedError(x: Any) = throw new IllegalArgumentException(s"Sign analysis meant to be run on programs without $x")
}

/**
 * Variant of [[tip.analysis.ForwardDependencies]] that includes call and return edges.
 */
trait SimpleInterproceduralDependencies extends Dependencies[GNode[AstNode]] {

  val cfg: ProgramCFG

  import cfg._

  override def dep(n: GNode[AstNode]) = n.succ.toSet ++ n.called ++ n.callersAfterCall

  override def invdep(n: GNode[AstNode]) = n.pred.toSet ++ n.calledExits ++ n.callers
}

/**
 * Common functionality for interprocedural analysis.
 */
trait InterprocSignAnalysisMisc {

  val lattice: MapLattice[GNode[AstNode], LiftLattice[MapLattice[AIdentifier, SignLattice.type]]]

  val cfg: ProgramCFG

  NormalisedCallsLanguageRestrictions.checkCfgLanguageRestriction(cfg)
  NoPointer.checkCfgLanguageRestriction(cfg)

  def evalArgs(formalParams: Seq[AIdentifier], actualParams: Seq[AExpr], state: lattice.l.l.Element): lattice.l.l.Element = {
    formalParams.zip(actualParams).foldLeft(lattice.l.l.bottom) { case (acc, (id, exp)) =>
      acc + (id -> lattice.l.l.l.eval(exp, state))
    }
  }
}

/**
 * Transfer functions for sign analysis (including interprocedural).
 * This version is for the basic worklist algorithm.
 */
trait InterprocSignAnalysisTransferFunctions extends MapLiftLatticeTransfer[GNode[AstNode]] with InterprocSignAnalysisMisc {

  override def transfer(n: GNode[AstNode], s: lattice.l.Element, o: lattice.Element): lattice.l.Element = {
    import cfg._
    import NormalisedCallsLanguageRestrictions._
    import lattice.l._

    n match {

      // call nodes (like no-ops here)
      case call: CallNode[AstNode] => join(n, s, o)

      // function entry nodes
      case funentry: FunEntry[AstNode] => {
        val fun = funentry.data.asInstanceOf[AFunDeclaration]
        val callSites = funentry.callers
        callSites.foldRight(s)((x, a) =>
          lattice.l.l.lub(a, evalArgs(fun.args, x.invocation.args, a))
        )
      }

      // function exit nodes (like no-ops here)
      case funexit: FunExit[AstNode] => join(n, s, o)

      // after-call nodes
      case aftercall: AfterCallNode[AstNode] => {
        val id = aftercall.assignment.leftId
        val idDef = id.meta.definition.get.asInstanceOf[AIdentifier]
        val vprime = o(aftercall.pred.head)
        val w = o(aftercall.calledExits.get)
        vprime + (idDef -> w(SignAnalysisCommons.returnId))
      }

      case GRealNode(_, _,_ , ret: AReturnStmt) => {
        val j = join(n, s, o)
        j + (SignAnalysisCommons.returnId -> lattice.l.l.l.eval(ret.value, j))
      }

      case _ => super.transfer(n, s, o)
    }
  }
}

/**
 * Transfer functions for sign analysis (including interprocedural), propagation style.
 * This is a variant of [[InterprocSignAnalysisTransferFunctions]] for use with [[tip.solvers.PropagationStrategy]].
 */
trait InterprocSignAnalysisTransferFunctionsWithPropagation extends MapLiftLatticeTransfer[GNode[AstNode]]
  with PropagationStrategy[GNode[AstNode]] with InterprocSignAnalysisMisc {

  override def transfer(n: GNode[AstNode], s: lattice.l.Element, o: lattice.Element): lattice.l.Element = {
    import cfg._
    import NormalisedCallsLanguageRestrictions._
    import lattice.l._

    // helper function that propagates dataflow from a function exit node to an after-call node
    def returnflow(funexit: FunExit[AstNode], aftercall: AfterCallNode[AstNode]) {
      val id = aftercall.assignment().leftId()
      propagate(o(aftercall.pred.head) + (id.meta.definition.get.asInstanceOf[AIdentifier] -> s(SignAnalysisCommons.returnId)), aftercall)
    }

    n match {

      // call nodes
      case call: CallNode[AstNode] => {
        val calledFun = call.called.get.data.asInstanceOf[AFunDeclaration]
        propagate(evalArgs(calledFun.args, call.invocation.args, s), call.called.head)
        returnflow(call.succ.head.calledExits.head, call.succ.head.asInstanceOf[AfterCallNode[AstNode]]) // make sure existing return flow gets propagated
        lattice.l.bottom // no flow directly to the after-call node
      }

      // function entry nodes (like no-op here)
      case funentry: FunEntry[AstNode] => s

      // function exit nodes
      case funexit: FunExit[AstNode] => {
        for (aftercall <- funexit.callersAfterCall) returnflow(funexit, aftercall)
        lattice.l.bottom // no successors for this kind of node, but we have to return something
      }

      // after-call nodes (like no-op here)
      case aftercall: AfterCallNode[AstNode] => s

      case GRealNode(_, _,_ , ret: AReturnStmt) => {
        s + (SignAnalysisCommons.returnId -> lattice.l.l.l.eval(ret.value, s))
      }

      case _ => super.transfer(n, s, o)
    }
  }
}

/**
 * Base class for intra-procedural sign analysis.
 */
abstract class IntraprocSignAnalysis(cfg: ControlFlowGraph[AstNode]) extends FlowSensitiveAnalysis(cfg)
  with IntraprocSignAnalysisTransferFunctions with MapLiftLatticeTransfer[GNode[AstNode]] with ForwardDependencies {

  /**
   * The lattice for the analysis, which is the lattice of maps associating with
   * each node a state, where the state is a map associating with each declared variable
   * an element of the sign lattice.
   */
  val lattice = new MapLattice(cfg.nodes, new LiftLattice(statelattice))
}

/**
 * Intraprocedural sign analysis that uses [[tip.solvers.SimpleFixpointSolver]].
 */
class IntraprocSignAnalysisSimpleSolver(cfg: ControlFlowGraph[AstNode])
  extends IntraprocSignAnalysis(cfg) with SimpleFixpointSolver with MapLatticeTransfer[GNode[AstNode]]

/**
 * Intraprocedural sign analysis that uses [[tip.solvers.WorklistFixpointSolver]].
 */
class IntraprocSignAnalysisWorklistSolver(cfg: ControlFlowGraph[AstNode])
  extends IntraprocSignAnalysis(cfg) with WorklistFixpointSolver[GNode[AstNode]]

/**
 * Intraprocedural sign analysis that uses [[tip.solvers.WorklistFixpointSolverWithInit]].
 */
class IntraprocSignAnalysisWorklistSolverWithInit(cfg: ControlFlowGraph[AstNode])
  extends IntraprocSignAnalysisWorklistSolver(cfg) with WorklistFixpointSolverWithInit[GNode[AstNode]] {

  val first = cfg.entries.values.toSet
}

/**
 * Intraprocedural sign analysis that uses [[tip.solvers.WorklistFixpointSolverWithInit]] with [[tip.solvers.PropagationStrategy]].
 */
class IntraprocSignAnalysisWorklistSolverWithInitAndPropagation(cfg: ControlFlowGraph[AstNode])
  extends IntraprocSignAnalysisWorklistSolverWithInit(cfg) with PropagationStrategy[GNode[AstNode]] {

  override val init = lattice.l.Lift(lattice.l.l.bottom)
}

/**
 * Interprocedural sign analysis that uses [[tip.solvers.WorklistFixpointSolverWithInit]] with the basic worklist algorithm.
 */
class InterprocSignAnalysisWorklistSolverWithInit(override val cfg: ProgramCFG)
  extends IntraprocSignAnalysisWorklistSolverWithInit(cfg) with InterprocSignAnalysisTransferFunctions
  with SimpleInterproceduralDependencies {

  override val first = Set(cfg.programEntry)
}

/**
 * Interprocedural sign analysis that uses [[tip.solvers.WorklistFixpointSolverWithInit]] with [[tip.solvers.PropagationStrategy]].
 */
class InterprocSignAnalysisWorklistSolverWithInitAndPropagation(override val cfg: ProgramCFG)
  extends IntraprocSignAnalysisWorklistSolverWithInitAndPropagation(cfg)
  with InterprocSignAnalysisTransferFunctionsWithPropagation with ForwardDependencies {

  override val first = Set(cfg.programEntry)
}