/**
 * @license
 * Copyright Google LLC All Rights Reserved.
 *
 * Use of this source code is governed by an MIT-style license that can be
 * found in the LICENSE file at https://angular.dev/license
 */

import * as o from '../../../../output/output_ast';
import {CONTEXT_NAME} from '../../../../render3/view/util';
import * as ir from '../../ir';
import {CompilationJob} from '../compilation';

/**
 * Optimize variables declared and used in the IR.
 *
 * Variables are eagerly generated by pipeline stages for all possible values that could be
 * referenced. This stage processes the list of declared variables and all variable usages,
 * and optimizes where possible. It performs 3 main optimizations:
 *
 *   * It transforms variable declarations to side effectful expressions when the
 *     variable is not used, but its initializer has global effects which other
 *     operations rely upon.
 *   * It removes variable declarations if those variables are not referenced and
 *     either they do not have global effects, or nothing relies on them.
 *   * It inlines variable declarations when those variables are only used once
 *     and the inlining is semantically safe.
 *
 * To guarantee correctness, analysis of "fences" in the instruction lists is used to determine
 * which optimizations are safe to perform.
 */
export function optimizeVariables(job: CompilationJob): void {
  for (const unit of job.units) {
    for (const expr of unit.functions) {
      inlineAlwaysInlineVariables(expr.ops);
    }
    inlineAlwaysInlineVariables(unit.create);
    inlineAlwaysInlineVariables(unit.update);

    for (const op of unit.create) {
      if (
        op.kind === ir.OpKind.Listener ||
        op.kind === ir.OpKind.Animation ||
        op.kind === ir.OpKind.AnimationListener ||
        op.kind === ir.OpKind.TwoWayListener
      ) {
        inlineAlwaysInlineVariables(op.handlerOps);
      } else if (op.kind === ir.OpKind.RepeaterCreate && op.trackByOps !== null) {
        inlineAlwaysInlineVariables(op.trackByOps);
      }
    }

    for (const expr of unit.functions) {
      optimizeVariablesInOpList(expr.ops, null);
      optimizeSaveRestoreView(expr.ops);
    }

    for (const op of unit.create) {
      if (
        op.kind === ir.OpKind.Listener ||
        op.kind === ir.OpKind.Animation ||
        op.kind === ir.OpKind.AnimationListener ||
        op.kind === ir.OpKind.TwoWayListener
      ) {
        optimizeVariablesInOpList(op.handlerOps, skipArrowFunctionOps);
        optimizeSaveRestoreView(op.handlerOps);
      } else if (op.kind === ir.OpKind.RepeaterCreate && op.trackByOps !== null) {
        optimizeVariablesInOpList(op.trackByOps, skipArrowFunctionOps);
      }
    }

    // Note that we skip over arrow function operations, because they are considered
    // separate boundaries that should not influence the surrounding create/update
    // operations. This is a side-effect of not being able to control which nested
    // ops `visitExpressionsInOp` will visit. Without this logic, variable references
    // inside the arrow function can throw off usage counting for things like view references.
    optimizeVariablesInOpList(unit.create, skipArrowFunctionOps);
    optimizeVariablesInOpList(unit.update, skipArrowFunctionOps);
  }
}

/**
 * A [fence](https://en.wikipedia.org/wiki/Memory_barrier) flag for an expression which indicates
 * how that expression can be optimized in relation to other expressions or instructions.
 *
 * `Fence`s are a bitfield, so multiple flags may be set on a single expression.
 */
enum Fence {
  /**
   * Empty flag (no fence exists).
   */
  None = 0b000,

  /**
   * A context read fence, meaning that the expression in question reads from the "current view"
   * context of the runtime.
   */
  ViewContextRead = 0b001,

  /**
   * A context write fence, meaning that the expression in question writes to the "current view"
   * context of the runtime.
   *
   * Note that all `ContextWrite` fences are implicitly `ContextRead` fences as operations which
   * change the view context do so based on the current one.
   */
  ViewContextWrite = 0b010,

  /**
   * Indicates that a call is required for its side-effects, even if nothing reads its result.
   *
   * This is also true of `ViewContextWrite` operations **if** they are followed by a
   * `ViewContextRead`.
   */
  SideEffectful = 0b100,
}

/**
 * Summary data collected for each `Op` in a list.
 *
 * Tracking this data per operation allows the optimizer to process operations at a higher level
 * than always scanning expressions.
 */
interface OpInfo {
  /**
   * A `Set` of variables referenced by expressions in this operation.
   */
  variablesUsed: Set<ir.XrefId>;

  /**
   * Flags indicating any `Fence`s present for this operation.
   */
  fences: Fence;
}

/** Predicate function that decides whether a certain expression should be processed. */
type ExpressionPredicate = (flags: ir.VisitorContextFlag) => boolean;

/** `ExpressionPredicate` that skips arrow function operations. */
function skipArrowFunctionOps(flags: ir.VisitorContextFlag): boolean {
  return !(flags & ir.VisitorContextFlag.InArrowFunctionOperation);
}

function inlineAlwaysInlineVariables(ops: ir.OpList<ir.CreateOp | ir.UpdateOp>): void {
  const vars = new Map<ir.XrefId, ir.VariableOp<ir.CreateOp | ir.UpdateOp>>();
  for (const op of ops) {
    if (op.kind === ir.OpKind.Variable && op.flags & ir.VariableFlags.AlwaysInline) {
      ir.visitExpressionsInOp(op, (expr) => {
        if (ir.isIrExpression(expr) && fencesForIrExpression(expr) !== Fence.None) {
          throw new Error(`AssertionError: A context-sensitive variable was marked AlwaysInline`);
        }
      });
      vars.set(op.xref, op);
    }

    ir.transformExpressionsInOp(
      op,
      (expr) => {
        if (expr instanceof ir.ReadVariableExpr && vars.has(expr.xref)) {
          const varOp = vars.get(expr.xref)!;
          // Inline by cloning, because we might inline into multiple places.
          return varOp.initializer.clone();
        }
        return expr;
      },
      ir.VisitorContextFlag.None,
    );
  }

  for (const op of vars.values()) {
    ir.OpList.remove(op as ir.CreateOp | ir.UpdateOp);
  }
}

/**
 * Process a list of operations and optimize variables within that list.
 */
function optimizeVariablesInOpList(
  ops: ir.OpList<ir.CreateOp | ir.UpdateOp>,
  predicate: ExpressionPredicate | null,
): void {
  const varDecls = new Map<ir.XrefId, ir.VariableOp<ir.CreateOp | ir.UpdateOp>>();
  const varUsages = new Map<ir.XrefId, number>();

  // Track variables that are used outside of the immediate operation list. For example, within
  // `ListenerOp` handler operations of listeners in the current operation list.
  const varRemoteUsages = new Set<ir.XrefId>();
  const opMap = new Map<ir.CreateOp | ir.UpdateOp, OpInfo>();

  // First, extract information about variables declared or used within the whole list.
  for (const op of ops) {
    if (op.kind === ir.OpKind.Variable) {
      if (varDecls.has(op.xref) || varUsages.has(op.xref)) {
        throw new Error(`Should not see two declarations of the same variable: ${op.xref}`);
      }
      varDecls.set(op.xref, op);
      varUsages.set(op.xref, 0);
    }

    opMap.set(op, collectOpInfo(op, predicate));
    countVariableUsages(op, varUsages, varRemoteUsages, predicate);
  }

  // The next step is to remove any variable declarations for variables that aren't used. The
  // variable initializer expressions may be side-effectful, so they may need to be retained as
  // expression statements.

  // Track whether we've seen an operation which reads from the view context yet. This is used to
  // determine whether a write to the view context in a variable initializer can be observed.
  let contextIsUsed = false;

  // Note that iteration through the list happens in reverse, which guarantees that we'll process
  // all reads of a variable prior to processing its declaration.
  for (const op of ops.reversed()) {
    const opInfo = opMap.get(op)!;

    if (op.kind === ir.OpKind.Variable && varUsages.get(op.xref)! === 0) {
      // This variable is unused and can be removed. We might need to keep the initializer around,
      // though, if something depends on it running.
      if (
        (contextIsUsed && opInfo.fences & Fence.ViewContextWrite) ||
        opInfo.fences & Fence.SideEffectful
      ) {
        // This variable initializer has a side effect which must be retained. Either:
        //  * it writes to the view context, and we know there is a future operation which depends
        //    on that write, or
        //  * it's an operation which is inherently side-effectful.
        // We can't remove the initializer, but we can remove the variable declaration itself and
        // replace it with a side-effectful statement.
        const stmtOp = ir.createStatementOp(op.initializer.toStmt()) as ir.UpdateOp;
        opMap.set(stmtOp, opInfo);
        ir.OpList.replace(op as ir.UpdateOp, stmtOp);
      } else {
        // It's safe to delete this entire variable declaration as nothing depends on it, even
        // side-effectfully. Note that doing this might make other variables unused. Since we're
        // iterating in reverse order, we should always be processing usages before declarations
        // and therefore by the time we get to a declaration, all removable usages will have been
        // removed.
        uncountVariableUsages(op, varUsages);
        ir.OpList.remove(op as ir.UpdateOp);
      }

      opMap.delete(op);
      varDecls.delete(op.xref);
      varUsages.delete(op.xref);
      continue;
    }

    // Does this operation depend on the view context?
    if (opInfo.fences & Fence.ViewContextRead) {
      contextIsUsed = true;
    }
  }

  // Next, inline any remaining variables with exactly one usage.
  const toInline: ir.XrefId[] = [];
  for (const [id, count] of varUsages) {
    const decl = varDecls.get(id)!;
    // We can inline variables that:
    //  - are used exactly once, and
    //  - are not used remotely
    // OR
    //  - are marked for always inlining
    const isAlwaysInline = !!(decl.flags & ir.VariableFlags.AlwaysInline);
    if (count !== 1 || isAlwaysInline) {
      // We can't inline this variable as it's used more than once.
      continue;
    }

    if (varRemoteUsages.has(id)) {
      // This variable is used once, but across an operation boundary, so it can't be inlined.
      continue;
    }

    toInline.push(id);
  }

  let candidate: ir.XrefId | undefined;
  while ((candidate = toInline.pop())) {
    // We will attempt to inline this variable. If inlining fails (due to fences for example),
    // no future operation will make inlining legal.
    const decl = varDecls.get(candidate)!;
    const varInfo = opMap.get(decl as ir.CreateOp | ir.UpdateOp)!;
    const isAlwaysInline = !!(decl.flags & ir.VariableFlags.AlwaysInline);

    if (isAlwaysInline) {
      throw new Error(
        `AssertionError: Found an 'AlwaysInline' variable after the always inlining pass.`,
      );
    }

    // Scan operations following the variable declaration and look for the point where that variable
    // is used. There should only be one usage given the precondition above.
    for (
      let targetOp = decl.next!;
      targetOp.kind !== ir.OpKind.ListEnd;
      targetOp = targetOp.next!
    ) {
      const opInfo = opMap.get(targetOp)!;

      // Is the variable used in this operation?
      if (opInfo.variablesUsed.has(candidate)) {
        if (!allowConservativeInlining(decl, targetOp)) {
          // We're in conservative mode, and this variable is not eligible for inlining into the
          // target operation in this mode.
          break;
        }

        // Yes, try to inline it. Inlining may not be successful if fences in this operation before
        // the variable's usage cannot be safely crossed.
        if (tryInlineVariableInitializer(candidate, decl.initializer, targetOp, varInfo.fences)) {
          // Inlining was successful! Update the tracking structures to reflect the inlined
          // variable.
          opInfo.variablesUsed.delete(candidate);

          // Add all variables used in the variable's initializer to its new usage site.
          for (const id of varInfo.variablesUsed) {
            opInfo.variablesUsed.add(id);
          }

          // Merge fences in the variable's initializer into its new usage site.
          opInfo.fences |= varInfo.fences;

          // Delete tracking info related to the declaration.
          varDecls.delete(candidate);
          varUsages.delete(candidate);
          opMap.delete(decl as ir.CreateOp | ir.UpdateOp);

          // And finally, delete the original declaration from the operation list.
          ir.OpList.remove(decl as ir.UpdateOp);
        }

        // Whether inlining succeeded or failed, we're done processing this variable.
        break;
      }

      // If the variable is not used in this operation, then we'd need to inline across it. Check if
      // that's safe to do.
      if (!safeToInlinePastFences(opInfo.fences, varInfo.fences)) {
        // We can't safely inline this variable beyond this operation, so don't proceed with
        // inlining this variable.
        break;
      }
    }
  }
}

/**
 * Given an `ir.Expression`, returns the `Fence` flags for that expression type.
 */
function fencesForIrExpression(expr: ir.Expression): Fence {
  switch (expr.kind) {
    case ir.ExpressionKind.NextContext:
      return Fence.ViewContextRead | Fence.ViewContextWrite;
    case ir.ExpressionKind.RestoreView:
      return Fence.ViewContextRead | Fence.ViewContextWrite | Fence.SideEffectful;
    case ir.ExpressionKind.StoreLet:
      return Fence.SideEffectful;
    case ir.ExpressionKind.Reference:
    case ir.ExpressionKind.ContextLetReference:
      return Fence.ViewContextRead;
    default:
      return Fence.None;
  }
}

/**
 * Build the `OpInfo` structure for the given `op`. This performs two operations:
 *
 *  * It tracks which variables are used in the operation's expressions.
 *  * It rolls up fence flags for expressions within the operation.
 */
function collectOpInfo(
  op: ir.CreateOp | ir.UpdateOp,
  predicate: ExpressionPredicate | null,
): OpInfo {
  let fences = Fence.None;
  const variablesUsed = new Set<ir.XrefId>();
  ir.visitExpressionsInOp(op, (expr, flags) => {
    if (!ir.isIrExpression(expr) || (predicate !== null && !predicate(flags))) {
      return;
    }

    if (expr.kind === ir.ExpressionKind.ReadVariable) {
      variablesUsed.add(expr.xref);
    } else {
      fences |= fencesForIrExpression(expr);
    }
  });
  return {fences, variablesUsed};
}

/**
 * Count the number of usages of each variable, being careful to track whether those usages are
 * local or remote.
 */
function countVariableUsages(
  op: ir.CreateOp | ir.UpdateOp,
  varUsages: Map<ir.XrefId, number>,
  varRemoteUsage: Set<ir.XrefId>,
  predicate: ExpressionPredicate | null,
): void {
  ir.visitExpressionsInOp(op, (expr, flags) => {
    if (!ir.isIrExpression(expr) || (predicate !== null && !predicate(flags))) {
      return;
    }

    if (expr.kind !== ir.ExpressionKind.ReadVariable) {
      return;
    }

    const count = varUsages.get(expr.xref);
    if (count === undefined) {
      // This variable is declared outside the current scope of optimization.
      return;
    }
    varUsages.set(expr.xref, count + 1);

    if (flags & ir.VisitorContextFlag.InChildOperation) {
      varRemoteUsage.add(expr.xref);
    }
  });
}

/**
 * Remove usages of a variable in `op` from the `varUsages` tracking.
 */
function uncountVariableUsages(
  op: ir.CreateOp | ir.UpdateOp,
  varUsages: Map<ir.XrefId, number>,
): void {
  ir.visitExpressionsInOp(op, (expr) => {
    if (!ir.isIrExpression(expr)) {
      return;
    }

    if (expr.kind !== ir.ExpressionKind.ReadVariable) {
      return;
    }

    const count = varUsages.get(expr.xref);
    if (count === undefined) {
      // This variable is declared outside the current scope of optimization.
      return;
    } else if (count === 0) {
      throw new Error(
        `Inaccurate variable count: ${expr.xref} - found another read but count is already 0`,
      );
    }
    varUsages.set(expr.xref, count - 1);
  });
}

/**
 * Checks whether it's safe to inline a variable across a particular operation.
 *
 * @param fences the fences of the operation which the inlining will cross
 * @param declFences the fences of the variable being inlined.
 */
function safeToInlinePastFences(fences: Fence, declFences: Fence): boolean {
  if (fences & Fence.ViewContextWrite) {
    // It's not safe to inline context reads across context writes.
    if (declFences & Fence.ViewContextRead) {
      return false;
    }
  } else if (fences & Fence.ViewContextRead) {
    // It's not safe to inline context writes across context reads.
    if (declFences & Fence.ViewContextWrite) {
      return false;
    }
  }
  return true;
}

/**
 * Attempt to inline the initializer of a variable into a target operation's expressions.
 *
 * This may or may not be safe to do. For example, the variable could be read following the
 * execution of an expression with fences that don't permit the variable to be inlined across them.
 */
function tryInlineVariableInitializer(
  id: ir.XrefId,
  initializer: o.Expression,
  target: ir.CreateOp | ir.UpdateOp,
  declFences: Fence,
): boolean {
  // We use `ir.transformExpressionsInOp` to walk the expressions and inline the variable if
  // possible. Since this operation is callback-based, once inlining succeeds or fails we can't
  // "stop" the expression processing, and have to keep track of whether inlining has succeeded or
  // is no longer allowed.
  let inlined = false;
  let inliningAllowed = true;

  ir.transformExpressionsInOp(
    target,
    (expr, flags) => {
      if (!ir.isIrExpression(expr)) {
        return expr;
      }

      if (inlined || !inliningAllowed) {
        // Either the inlining has already succeeded, or we've passed a fence that disallows inlining
        // at this point, so don't try.
        return expr;
      } else if (
        flags & ir.VisitorContextFlag.InChildOperation &&
        declFences & Fence.ViewContextRead
      ) {
        // We cannot inline variables that are sensitive to the current context across operation
        // boundaries.
        return expr;
      }

      switch (expr.kind) {
        case ir.ExpressionKind.ReadVariable:
          if (expr.xref === id) {
            // This is the usage site of the variable. Since nothing has disallowed inlining, it's
            // safe to inline the initializer here.
            inlined = true;
            return initializer;
          }
          break;
        default:
          // For other types of `ir.Expression`s, whether inlining is allowed depends on their fences.
          const exprFences = fencesForIrExpression(expr);
          inliningAllowed = inliningAllowed && safeToInlinePastFences(exprFences, declFences);
          break;
      }
      return expr;
    },
    ir.VisitorContextFlag.None,
  );
  return inlined;
}

/**
 * Determines whether inlining of `decl` should be allowed in "conservative" mode.
 *
 * In conservative mode, inlining behavior is limited to those operations which the
 * `TemplateDefinitionBuilder` supported, with the goal of producing equivalent output.
 */
function allowConservativeInlining(
  decl: ir.VariableOp<ir.CreateOp | ir.UpdateOp>,
  target: ir.Op<ir.CreateOp | ir.UpdateOp>,
): boolean {
  // TODO(alxhub): understand exactly how TemplateDefinitionBuilder approaches inlining, and record
  // that behavior here.
  switch (decl.variable.kind) {
    case ir.SemanticVariableKind.Identifier:
      if (decl.initializer instanceof o.ReadVarExpr && decl.initializer.name === CONTEXT_NAME) {
        // Although TemplateDefinitionBuilder is cautious about inlining, we still want to do so
        // when the variable is the context, to imitate its behavior with aliases in control flow
        // blocks. This quirky behavior will become dead code once compatibility mode is no longer
        // supported.
        return true;
      }
      return false;
    case ir.SemanticVariableKind.Context:
      // Context can only be inlined into other variables.
      return target.kind === ir.OpKind.Variable;
    default:
      return true;
  }
}

/**
 * After variables have been optimized in nested ops (e.g. handlers or functions), we may end up
 * with `saveView`/`restoreView` calls that aren't necessary since all the references to the view
 * were optimized away. This function removes the ops related to the view restoration.
 */
function optimizeSaveRestoreView(ops: ir.OpList<ir.UpdateOp>): void {
  const head = ops.head.next;
  const tail = ops.tail.prev;

  // We can only optimize if we have two ops:
  // 1. A call to `restoreView`.
  // 2. A return statement with a `resetView` in it.
  if (
    head !== null &&
    tail !== null &&
    head.next === tail &&
    head.kind === ir.OpKind.Statement &&
    head.statement instanceof o.ExpressionStatement &&
    head.statement.expr instanceof ir.RestoreViewExpr &&
    tail.kind === ir.OpKind.Statement &&
    tail.statement instanceof o.ReturnStatement &&
    tail.statement.value instanceof ir.ResetViewExpr
  ) {
    ir.OpList.remove<ir.UpdateOp>(head);
    tail.statement.value = tail.statement.value.expr;
  }
}
