/**
 * 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 {
  ReactiveFunction,
  ReactiveScopeBlock,
  ReactiveStatement,
  ReactiveTerminal,
  ReactiveTerminalStatement,
} from "../HIR/HIR";
import { assertExhaustive } from "../Utils/utils";
import {
  ReactiveFunctionTransform,
  Transformed,
  visitReactiveFunction,
} from "./visitors";

/*
 * Given a reactive function, flattens any scopes contained within a loop construct.
 * We won't initially support memoization within loops though this is possible in the future.
 */
export function flattenReactiveLoops(fn: ReactiveFunction): void {
  visitReactiveFunction(fn, new Transform(), false);
}

class Transform extends ReactiveFunctionTransform<boolean> {
  override transformScope(
    scope: ReactiveScopeBlock,
    isWithinLoop: boolean
  ): Transformed<ReactiveStatement> {
    this.visitScope(scope, isWithinLoop);
    if (isWithinLoop) {
      return {
        kind: "replace",
        value: {
          kind: "pruned-scope",
          scope: scope.scope,
          instructions: scope.instructions,
        },
      };
    } else {
      return { kind: "keep" };
    }
  }

  override visitTerminal(
    stmt: ReactiveTerminalStatement<ReactiveTerminal>,
    isWithinLoop: boolean
  ): void {
    switch (stmt.terminal.kind) {
      // Loop terminals flatten nested scopes
      case "do-while":
      case "while":
      case "for":
      case "for-of":
      case "for-in": {
        this.traverseTerminal(stmt, true);
        break;
      }
      // Non-loop terminals passthrough is contextual, inherits the parent isWithinScope
      case "try":
      case "label":
      case "break":
      case "continue":
      case "if":
      case "return":
      case "switch":
      case "throw": {
        this.traverseTerminal(stmt, isWithinLoop);
        break;
      }
      default: {
        assertExhaustive(
          stmt.terminal,
          `Unexpected terminal kind \`${(stmt.terminal as any).kind}\``
        );
      }
    }
  }
}