/**
 * 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,
  ReactiveTerminalStatement,
} from '../HIR/HIR';
import {
  ReactiveFunctionTransform,
  Transformed,
  visitReactiveFunction,
} from './visitors';

// Converts scopes without outputs into regular blocks.
export function pruneUnusedScopes(fn: ReactiveFunction): void {
  visitReactiveFunction(fn, new Transform(), {
    hasReturnStatement: false,
  } as State);
}

type State = {
  hasReturnStatement: boolean;
};

class Transform extends ReactiveFunctionTransform<State> {
  override visitTerminal(stmt: ReactiveTerminalStatement, state: State): void {
    this.traverseTerminal(stmt, state);
    if (stmt.terminal.kind === 'return') {
      state.hasReturnStatement = true;
    }
  }
  override transformScope(
    scopeBlock: ReactiveScopeBlock,
    _state: State,
  ): Transformed<ReactiveStatement> {
    const scopeState: State = {hasReturnStatement: false};
    this.visitScope(scopeBlock, scopeState);
    if (
      !scopeState.hasReturnStatement &&
      scopeBlock.scope.reassignments.size === 0 &&
      (scopeBlock.scope.declarations.size === 0 ||
        /*
         * Can prune scopes where all declarations bubbled up from inner
         * scopes
         */
        !hasOwnDeclaration(scopeBlock))
    ) {
      return {
        kind: 'replace',
        value: {
          kind: 'pruned-scope',
          scope: scopeBlock.scope,
          instructions: scopeBlock.instructions,
        },
      };
    } else {
      return {kind: 'keep'};
    }
  }
}

/*
 * Does the scope block declare any values of its own? This can return
 * false if all the block's declarations are propagated from nested scopes.
 */
function hasOwnDeclaration(block: ReactiveScopeBlock): boolean {
  for (const declaration of block.scope.declarations.values()) {
    if (declaration.scope.id === block.scope.id) {
      return true;
    }
  }
  return false;
}