/**
 * 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 { CompilerError } from "..";
import { HIRFunction, IdentifierId, Place } from "../HIR";
import { printPlace } from "../HIR/PrintHIR";
import {
  eachInstructionValueLValue,
  eachPatternOperand,
} from "../HIR/visitors";

/**
 * Validates that all store/load references to a given named identifier align with the
 * "kind" of that variable (normal variable or context variable). For example, a context
 * variable may not be loaded/stored with regular StoreLocal/LoadLocal/Destructure instructions.
 */
export function validateContextVariableLValues(fn: HIRFunction): void {
  const identifierKinds: IdentifierKinds = new Map();
  validateContextVariableLValuesImpl(fn, identifierKinds);
}

function validateContextVariableLValuesImpl(
  fn: HIRFunction,
  identifierKinds: IdentifierKinds
): void {
  for (const [, block] of fn.body.blocks) {
    for (const instr of block.instructions) {
      const { value } = instr;
      switch (value.kind) {
        case "DeclareContext":
        case "StoreContext": {
          visit(identifierKinds, value.lvalue.place, "context");
          break;
        }
        case "LoadContext": {
          visit(identifierKinds, value.place, "context");
          break;
        }
        case "StoreLocal":
        case "DeclareLocal": {
          visit(identifierKinds, value.lvalue.place, "local");
          break;
        }
        case "LoadLocal": {
          visit(identifierKinds, value.place, "local");
          break;
        }
        case "PostfixUpdate":
        case "PrefixUpdate": {
          visit(identifierKinds, value.lvalue, "local");
          break;
        }
        case "Destructure": {
          for (const lvalue of eachPatternOperand(value.lvalue.pattern)) {
            visit(identifierKinds, lvalue, "destructure");
          }
          break;
        }
        case "ObjectMethod":
        case "FunctionExpression": {
          validateContextVariableLValuesImpl(
            value.loweredFunc.func,
            identifierKinds
          );
          break;
        }
        default: {
          for (const _ of eachInstructionValueLValue(value)) {
            CompilerError.throwTodo({
              reason:
                "ValidateContextVariableLValues: unhandled instruction variant",
              loc: value.loc,
              description: `Handle '${value.kind} lvalues`,
              suggestions: null,
            });
          }
        }
      }
    }
  }
}

type IdentifierKinds = Map<
  IdentifierId,
  { place: Place; kind: "local" | "context" | "destructure" }
>;

function visit(
  identifiers: IdentifierKinds,
  place: Place,
  kind: "local" | "context" | "destructure"
): void {
  const prev = identifiers.get(place.identifier.id);
  if (prev !== undefined) {
    const wasContext = prev.kind === "context";
    const isContext = kind === "context";
    if (wasContext !== isContext) {
      if (prev.kind === "destructure" || kind === "destructure") {
        CompilerError.throwTodo({
          reason: `Support destructuring of context variables`,
          loc: kind === "destructure" ? place.loc : prev.place.loc,
          description: null,
          suggestions: null,
        });
      }

      CompilerError.invariant(false, {
        reason: `Expected all references to a variable to be consistently local or context references`,
        loc: place.loc,
        description: `Identifier ${printPlace(
          place
        )} is referenced as a ${kind} variable, but was previously referenced as a ${prev} variable`,
        suggestions: null,
      });
    }
  }
  identifiers.set(place.identifier.id, { place, kind });
}