/**
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

import {visitReactiveFunction} from '.';
import {Effect} from '..';
import {
  Environment,
  GeneratedSource,
  InstructionKind,
  ReactiveFunction,
  ReactiveScope,
  ReactiveScopeBlock,
  ReactiveStatement,
  ReactiveTerminalStatement,
  makeInstructionId,
  promoteTemporary,
} from '../HIR';
import {createTemporaryPlace} from '../HIR/HIRBuilder';
import {EARLY_RETURN_SENTINEL} from './CodegenReactiveFunction';
import {ReactiveFunctionTransform, Transformed} from './visitors';

/**
 * This pass ensures that reactive blocks honor the control flow behavior of the
 * original code including early return semantics. Specifically, if a reactive
 * scope early returned during the previous execution and the inputs to that block
 * have not changed, then the code should early return (with the same value) again.
 *
 * Example:
 *
 * ```javascript
 * let x = [];
 * if (props.cond) {
 *   x.push(12);
 *   return x;
 * } else {
 *   return foo();
 * }
 * ```
 *
 * Imagine that this code is called twice in a row with props.cond = true. Both
 * times it should return the same object (===), an array `[12]`.
 *
 * The compilation strategy is as follows. For each top-level reactive scope
 * that contains (transitively) an early return:
 *
 * - Label the scope
 * - Synthesize a new temporary, eg `t0`, and set it as a declaration of the scope.
 *   This will represent the possibly-unset return value for that scope.
 * - Make the first instruction of the scope the declaration of that temporary,
 *   assigning a sentinel value (can reuse the same symbol as we use for cache slots).
 *   This assignment ensures that if we don't take an early return, that the value
 *   is the sentinel.
 * - Replace all `return` statements with:
 *   - An assignment of the temporary with the value being returned.
 *   - A `break` to the reactive scope's label.
 *
 * Finally, CodegenReactiveScope adds an if check following the reactive scope:
 * if the early return temporary value is *not* the sentinel value, we early return
 * it. Otherwise, execution continues.
 *
 * For the above example that looks roughly like:
 *
 * ```
 * let t0;
 * if (props.cond !== $[0]) {
 *   t0 = Symbol.for('react.memo_cache_sentinel');
 *   bb0: {
 *     let x = [];
 *     if (props.cond) {
 *       x.push(12);
 *       t0 = x;
 *       break bb0;
 *     } else {
 *       let t1;
 *       if ($[1] === Symbol.for('react.memo_cache_sentinel')) {
 *         t1 = foo();
 *         $[1] = t1;
 *       } else {
 *         t1 = $[1];
 *       }
 *       t0 = t1;
 *       break bb0;
 *     }
 *   }
 *   $[0] = props.cond;
 *   $[2] = t0;
 * } else {
 *   t0 = $[2];
 * }
 * // This part added in CodegenReactiveScope:
 * if (t0 !== Symbol.for('react.memo_cache_sentinel')) {
 *   return t0;
 * }
 * ```
 */
export function propagateEarlyReturns(fn: ReactiveFunction): void {
  visitReactiveFunction(fn, new Transform(fn.env), {
    withinReactiveScope: false,
    earlyReturnValue: null,
  });
}

type State = {
  /**
   * Are we within a reactive scope? We use this for two things:
   * - When we find an early return, transform it to an assign+break
   *   only if we're in a reactive scope
   * - Annotate reactive scopes that contain early returns...but only
   *   the outermost reactive scope, we can't do this for nested
   *   scopes.
   */
  withinReactiveScope: boolean;

  /**
   * Store early return information to bubble it back up to the outermost
   * reactive scope
   */
  earlyReturnValue: ReactiveScope['earlyReturnValue'];
};

class Transform extends ReactiveFunctionTransform<State> {
  env: Environment;
  constructor(env: Environment) {
    super();
    this.env = env;
  }

  override visitScope(
    scopeBlock: ReactiveScopeBlock,
    parentState: State,
  ): void {
    /**
     * Exit early if an earlier pass has already created an early return,
     * which may happen in alternate compiler configurations.
     */
    if (scopeBlock.scope.earlyReturnValue !== null) {
      return;
    }

    const innerState: State = {
      withinReactiveScope: true,
      earlyReturnValue: parentState.earlyReturnValue,
    };
    this.traverseScope(scopeBlock, innerState);

    const earlyReturnValue = innerState.earlyReturnValue;
    if (earlyReturnValue !== null) {
      if (!parentState.withinReactiveScope) {
        // This is the outermost scope wrapping an early return, store the early return information
        scopeBlock.scope.earlyReturnValue = earlyReturnValue;
        scopeBlock.scope.declarations.set(earlyReturnValue.value.id, {
          identifier: earlyReturnValue.value,
          scope: scopeBlock.scope,
        });

        const instructions = scopeBlock.instructions;
        const loc = earlyReturnValue.loc;
        const sentinelTemp = createTemporaryPlace(this.env, loc);
        const symbolTemp = createTemporaryPlace(this.env, loc);
        const forTemp = createTemporaryPlace(this.env, loc);
        const argTemp = createTemporaryPlace(this.env, loc);
        scopeBlock.instructions = [
          {
            kind: 'instruction',
            instruction: {
              id: makeInstructionId(0),
              loc,
              lvalue: {...symbolTemp},
              value: {
                kind: 'LoadGlobal',
                binding: {
                  kind: 'Global',
                  name: 'Symbol',
                },
                loc,
              },
            },
          },
          {
            kind: 'instruction',
            instruction: {
              id: makeInstructionId(0),
              loc,
              lvalue: {...forTemp},
              value: {
                kind: 'PropertyLoad',
                object: {...symbolTemp},
                property: 'for',
                loc,
              },
            },
          },
          {
            kind: 'instruction',
            instruction: {
              id: makeInstructionId(0),
              loc,
              lvalue: {...argTemp},
              value: {
                kind: 'Primitive',
                value: EARLY_RETURN_SENTINEL,
                loc,
              },
            },
          },
          {
            kind: 'instruction',
            instruction: {
              id: makeInstructionId(0),
              loc,
              lvalue: {...sentinelTemp},
              value: {
                kind: 'MethodCall',
                receiver: symbolTemp,
                property: forTemp,
                args: [argTemp],
                loc,
              },
            },
          },
          {
            kind: 'instruction',
            instruction: {
              id: makeInstructionId(0),
              loc,
              lvalue: null,
              value: {
                kind: 'StoreLocal',
                loc,
                type: null,
                lvalue: {
                  kind: InstructionKind.Let,
                  place: {
                    kind: 'Identifier',
                    effect: Effect.ConditionallyMutate,
                    loc,
                    reactive: true,
                    identifier: earlyReturnValue.value,
                  },
                },
                value: {...sentinelTemp},
              },
            },
          },
          {
            kind: 'terminal',
            label: {
              id: earlyReturnValue.label,
              implicit: false,
            },
            terminal: {
              kind: 'label',
              id: makeInstructionId(0),
              loc: GeneratedSource,
              block: instructions,
            },
          },
        ];
      } else {
        /*
         * Not the outermost scope, but we save the early return information in case there are other
         * early returns within the same outermost scope
         */
        parentState.earlyReturnValue = earlyReturnValue;
      }
    }
  }

  override transformTerminal(
    stmt: ReactiveTerminalStatement,
    state: State,
  ): Transformed<ReactiveStatement> {
    if (state.withinReactiveScope && stmt.terminal.kind === 'return') {
      const loc = stmt.terminal.value.loc;
      let earlyReturnValue: ReactiveScope['earlyReturnValue'];
      if (state.earlyReturnValue !== null) {
        earlyReturnValue = state.earlyReturnValue;
      } else {
        const identifier = createTemporaryPlace(this.env, loc).identifier;
        promoteTemporary(identifier);
        earlyReturnValue = {
          label: this.env.nextBlockId,
          loc,
          value: identifier,
        };
      }
      state.earlyReturnValue = earlyReturnValue;
      return {
        kind: 'replace-many',
        value: [
          {
            kind: 'instruction',
            instruction: {
              id: makeInstructionId(0),
              loc,
              lvalue: null,
              value: {
                kind: 'StoreLocal',
                loc,
                type: null,
                lvalue: {
                  kind: InstructionKind.Reassign,
                  place: {
                    kind: 'Identifier',
                    identifier: earlyReturnValue.value,
                    effect: Effect.Capture,
                    loc,
                    reactive: true,
                  },
                },
                value: stmt.terminal.value,
              },
            },
          },
          {
            kind: 'terminal',
            label: null,
            terminal: {
              kind: 'break',
              id: makeInstructionId(0),
              loc,
              targetKind: 'labeled',
              target: earlyReturnValue.label,
            },
          },
        ],
      };
    }
    this.traverseTerminal(stmt, state);
    return {kind: 'keep'};
  }
}
