{-# LANGUAGE OverloadedStrings #-}

module Test.Arm64.Suite
  ( tests,
  )
where

import Aihc.Arm64
  ( Arm64Error (..),
    ObservedProgram (..),
    buildLinkLayout,
    compileModule,
    compileObservedFunction,
    compileProgram,
    snapshotSourcePath,
    targetTriple,
    validateProgramPrimitives,
  )
import Aihc.Grin
import Aihc.Native
  ( NativeTarget (AppleArm64),
    RuntimeGarbageCollector (..),
    RuntimePlan (..),
    runtimePlan,
  )
import Aihc.Tc (Levity (..), RuntimeRep (..), Unique (..))
import Aihc.Testing.EvalFixture (EvalCase (..), compileEvalCase, evalBindingNameInProgram, loadEvalCases)
import Aihc.Testing.ExceptionProgram (synchronousExceptionProgram)
import Aihc.Testing.SchedulerProgram (blackholeSchedulerProgram, schedulerProgram, stdioSchedulerProgram)
import Control.Concurrent (threadDelay)
import Control.Exception (bracket)
import Control.Monad (when)
import Data.Aeson (FromJSON (..), withObject, (.:), (.:?))
import Data.List (find, isInfixOf)
import Data.Map.Strict qualified as Map
import Data.Maybe (fromMaybe)
import Data.Text qualified as T
import Data.Text.IO qualified as TIO
import Data.Word (Word64)
import Data.Yaml qualified as Y
import System.Directory (createDirectory, doesDirectoryExist, getCurrentDirectory, getTemporaryDirectory, removeDirectoryRecursive, removeFile)
import System.Exit (ExitCode (..))
import System.FilePath (takeDirectory, (</>))
import System.IO (hClose, hFlush, hPutStr, openTempFile)
import System.Info (arch, os)
import System.Process (CreateProcess (..), StdStream (..), createProcess, proc, readProcessWithExitCode, waitForProcess)
import Test.Tasty (TestTree, testGroup)
import Test.Tasty.HUnit (assertBool, assertEqual, assertFailure, testCase)

tests :: TestTree
tests =
  testGroup
    "aihc-arm64"
    [ testCase "rejects unresolved representation-polymorphic native layouts" $ do
        let runtimeRep = RuntimeRepVar (Unique 1)
            program =
              GrinProgram
                { grinConstructors = [("Box", [[runtimeRep]])],
                  grinPrimitives = [],
                  grinForeignCalls = [],
                  grinExternalGlobals = [],
                  grinExternalFunctions = [],
                  grinWhnfGlobals = [],
                  grinCafs = [],
                  grinFunctions = []
                }
        assertEqual
          "native representation diagnostic"
          (Left (Arm64UnsupportedRuntimeRep runtimeRep))
          (compileProgram "missing" (expectGcGrin program)),
      testCase "keeps unsupported dormant primitives out of linked programs" $ do
        let primitive = GrinVar "unsupported#" 1 (BoxedRep Lifted)
            program =
              GrinProgram
                { grinConstructors = [],
                  grinPrimitives = [(primitive, 2)],
                  grinForeignCalls = [],
                  grinExternalGlobals = [],
                  grinExternalFunctions = [],
                  grinWhnfGlobals = [],
                  grinCafs = [],
                  grinFunctions = []
                }
        assertEqual "linked primitive validation" (Left (Arm64UnsupportedPrimitive "unsupported#")) (validateProgramPrimitives program)
        case compileModule (buildLinkLayout [program]) "_aihc_init_test" (expectGcGrin program) of
          Left err -> assertFailure ("relocatable module rejected a dormant primitive: " <> show err)
          Right assembly -> assertBool "module initializer" (".globl _aihc_init_test" `T.isInfixOf` assembly),
      testCase "adds Int# values with wrapping machine arithmetic" $ do
        let entryName = FunctionName "int_add"
            result = GrinVar "result" 2 IntRep
            program =
              GrinProgram
                { grinConstructors = [],
                  grinPrimitives = [(GrinVar "+#" 1 IntRep, 2)],
                  grinForeignCalls = [],
                  grinExternalGlobals = [],
                  grinExternalFunctions = [],
                  grinWhnfGlobals = [],
                  grinCafs = [],
                  grinFunctions =
                    [ GrinFunction
                        { grinFunctionName = entryName,
                          grinFunctionLinkName = Nothing,
                          grinFunctionParameters = [],
                          grinFunctionResultRep = IntRep,
                          grinFunctionBody =
                            GrinBind
                              [result]
                              ( GrinPrimitiveCall
                                  IntRep
                                  "+#"
                                  [ GrinLitValue (GrinLitInt IntRep 9223372036854775807),
                                    GrinLitValue (GrinLitInt IntRep 1)
                                  ]
                              )
                              (GrinConstant [GrinVarValue result])
                        }
                    ]
                }
        assertEqual "direct GRIN lint" [] (lintProgram program)
        assertEqual "linked primitive validation" (Right ()) (validateProgramPrimitives program)
        let gc = expectGcGrin program
        assertEqual "GC-GRIN lint" [] (lintProgram (gcGrinProgram gc))
        observed <-
          case compileObservedFunction entryName gc of
            Left err -> assertFailure ("native Int# addition compilation failed: " <> show err)
            Right value -> pure value
        assertBool "emits a wrapping 64-bit add" ("  add x0, x9, x0" `T.isInfixOf` observedAssembly observed)
        when (arch == "aarch64" && os == "darwin") $ do
          native <- runObservedProgram observed
          assertEqual "native result" (Right "return: -9223372036854775808\nheap: []\nallocations: 0\n") native,
      testCase "multiplies Word# values to a high/low pair" $ do
        let entryName = FunctionName "word_mul_wide"
            high = GrinVar "high" 2 WordRep
            low = GrinVar "low" 3 WordRep
            wordLiteral = GrinLitValue . GrinLitInt WordRep
            program =
              GrinProgram
                { grinConstructors = [],
                  grinPrimitives = [(GrinVar "timesWord2#" 1 (TupleRep [WordRep, WordRep]), 2)],
                  grinForeignCalls = [],
                  grinExternalGlobals = [],
                  grinExternalFunctions = [],
                  grinWhnfGlobals = [],
                  grinCafs = [],
                  grinFunctions =
                    [ GrinFunction
                        { grinFunctionName = entryName,
                          grinFunctionLinkName = Nothing,
                          grinFunctionParameters = [],
                          grinFunctionResultRep = TupleRep [WordRep, WordRep],
                          grinFunctionBody =
                            GrinBind
                              [high, low]
                              (GrinPrimitiveCall (TupleRep [WordRep, WordRep]) "timesWord2#" [wordLiteral 0xffffffffffffffff, wordLiteral 2])
                              (GrinConstant [GrinVarValue high, GrinVarValue low])
                        }
                    ]
                }
        observed <-
          case compileObservedFunction entryName (expectGcGrin program) of
            Left err -> assertFailure ("native timesWord2# compilation failed: " <> show err)
            Right value -> pure value
        assertBool "emits high-half multiplication" ("  umulh x10, x9, x0" `T.isInfixOf` observedAssembly observed)
        assertBool "emits low-half multiplication" ("  mul x11, x9, x0" `T.isInfixOf` observedAssembly observed),
      testCase "canonicalizes narrow signed literals in machine-word slots" $ do
        let functionName = FunctionName "narrow_code"
            program =
              GrinProgram
                { grinConstructors = [],
                  grinPrimitives = [],
                  grinForeignCalls = [],
                  grinExternalGlobals = [],
                  grinExternalFunctions = [],
                  grinWhnfGlobals = [],
                  grinCafs = [],
                  grinFunctions =
                    [ GrinFunction
                        { grinFunctionName = functionName,
                          grinFunctionLinkName = Nothing,
                          grinFunctionParameters = [],
                          grinFunctionResultRep = Int8Rep,
                          grinFunctionBody = GrinConstant [GrinLitValue (GrinLitInt Int8Rep 255)]
                        }
                    ]
                }
        case compileModule (buildLinkLayout [program]) "_aihc_init_narrow" (expectGcGrin program) of
          Left err -> assertFailure ("native compilation failed: " <> show err)
          Right assembly -> assertBool "255 :: Int8# is stored as -1" ("ldr x0, =-1" `T.isInfixOf` assembly),
      testCase "passes static Addr# literals to native foreign calls" $ do
        let functionName = FunctionName "puts_addr"
            foreignCall =
              GrinForeignCall
                { grinForeignCallName = "$ffi$puts",
                  grinForeignCallSymbol = "puts",
                  grinForeignCallSignature =
                    GrinForeignSignature
                      { grinForeignArgumentTypes = [GrinForeignAddr],
                        grinForeignResultType = GrinForeignInt32,
                        grinForeignEffect = GrinForeignPure
                      }
                }
            program =
              GrinProgram
                { grinConstructors = [],
                  grinPrimitives = [],
                  grinForeignCalls = [foreignCall],
                  grinExternalGlobals = [],
                  grinExternalFunctions = [],
                  grinWhnfGlobals = [],
                  grinCafs = [],
                  grinFunctions =
                    [ GrinFunction
                        { grinFunctionName = functionName,
                          grinFunctionLinkName = Nothing,
                          grinFunctionParameters = [],
                          grinFunctionResultRep = Int32Rep,
                          grinFunctionBody = GrinForeignCallExpr foreignCall [GrinLitValue (GrinLitAddr "\xFF\0bar")]
                        }
                    ]
                }
        case compileModule (buildLinkLayout [program]) "_aihc_init_addr" (expectGcGrin program) of
          Left err -> assertFailure ("native compilation failed: " <> show err)
          Right assembly -> do
            assertBool "loads the static string address" ("adrp x0, .Laihc_addr_0@PAGE" `T.isInfixOf` assembly)
            assertBool "emits NUL-terminated Latin-1" (".byte 255, 0, 98, 97, 114, 0" `T.isInfixOf` assembly)
            assertBool "calls puts" ("bl _puts" `T.isInfixOf` assembly),
      testCase "returns unboxed tuples as direct machine values" $ do
        let functionName = FunctionName "pair_code"
            program =
              GrinProgram
                { grinConstructors = [],
                  grinPrimitives = [],
                  grinForeignCalls = [],
                  grinExternalGlobals = [],
                  grinExternalFunctions = [],
                  grinWhnfGlobals = [],
                  grinCafs = [],
                  grinFunctions =
                    [ GrinFunction
                        { grinFunctionName = functionName,
                          grinFunctionLinkName = Nothing,
                          grinFunctionParameters = [],
                          grinFunctionResultRep = TupleRep [TupleRep [], IntRep, WordRep],
                          grinFunctionBody =
                            GrinConstant
                              [ GrinLitValue (GrinLitInt IntRep 1),
                                GrinLitValue (GrinLitInt WordRep 2)
                              ]
                        }
                    ]
                }
        case compileModule (buildLinkLayout [program]) "_aihc_init_pair" (expectGcGrin program) of
          Left err -> assertFailure ("native compilation failed: " <> show err)
          Right assembly -> do
            assertBool "passes two values" ("ldr x2, =2" `T.isInfixOf` assembly)
            assertBool "enters the continuation through registers" ("b .Laihc_enter" `T.isInfixOf` assembly)
            assertBool "does not call a C continuation adapter" (not ("_aihc_continue_values" `T.isInfixOf` assembly)),
      testCase "exports stable entries and branches directly to dependency code" $ do
        let identityName = FunctionName "$entry$identity"
            callerName = FunctionName "$entry$caller"
            argument = GrinVar "argument" 1 (BoxedRep Lifted)
            program =
              GrinProgram
                { grinConstructors = [],
                  grinPrimitives = [],
                  grinForeignCalls = [],
                  grinExternalGlobals = [],
                  grinExternalFunctions =
                    [ GrinCodeInfo
                        { grinCodeSourceName = "identity",
                          grinCodeFunctionName = identityName,
                          grinCodeParameterLayouts = [[BoxedRep Lifted]],
                          grinCodeResultRep = BoxedRep Lifted
                        }
                    ],
                  grinWhnfGlobals = [],
                  grinCafs = [],
                  grinFunctions =
                    [ GrinFunction
                        { grinFunctionName = callerName,
                          grinFunctionLinkName = Just "caller",
                          grinFunctionParameters = [argument],
                          grinFunctionResultRep = BoxedRep Lifted,
                          grinFunctionBody = GrinCall (BoxedRep Lifted) identityName [GrinVarValue argument]
                        }
                    ]
                }
        case compileModule (buildLinkLayout [program]) "_aihc_init_direct_call" (expectGcGrin program) of
          Left err -> assertFailure ("native compilation failed: " <> show err)
          Right assembly -> do
            assertBool "exports caller entry" (".globl _aihc_entry_caller" `T.isInfixOf` assembly)
            assertBool "branches to dependency entry" ("b _aihc_entry_identity" `T.isInfixOf` assembly)
            assertBool "does not publish direct-call arguments through the machine" (not ("str x8, [x22, #0]" `T.isInfixOf` assembly)),
      testGroup "raw GRIN heap snapshots" (map snapshotTest snapshotCases),
      testCase "case and apply never evaluate operands implicitly" $ do
        case compileModule (buildLinkLayout [explicitEvaluationProgram]) "_aihc_init_explicit_eval" (expectGcGrin explicitEvaluationProgram) of
          Left err -> assertFailure ("native compilation failed: " <> show err)
          Right assembly ->
            assertBool "generated case and apply contain no direct-style eval call" (not ("bl _aihc_eval\n" `T.isInfixOf` assembly))
        pure (),
      testCase "dynamic CPS transfers branch to runtime-selected entries" $ do
        case compileModule (buildLinkLayout [explicitEvaluationProgram]) "_aihc_init_tail_dispatch" (expectGcGrin explicitEvaluationProgram) of
          Left err -> assertFailure ("native compilation failed: " <> show err)
          Right assembly -> do
            assertBool
              "slow apply returns a value that is passed to the continuation in registers"
              ("bl _aihc_apply_slow" `T.isInfixOf` assembly && "b .Laihc_enter" `T.isInfixOf` assembly)
            assertBool
              "generated code does not reload a scheduled entry"
              (not ("ldr x9, [x22, #0]\n  br x9" `T.isInfixOf` assembly))
        pure (),
      testCase "runtime object ABI compiles cleanly on the host C compiler" $
        withTempDirectory "aihc-arm64-runtime" $ \directory -> do
          runtimeArguments <- nativeRuntimeArguments RuntimeGcCalloc
          snapshotRuntime <- snapshotSourcePath
          let executable = directory </> "runtime-check"
          (compilerExit, _compilerOut, compilerErr) <-
            readProcessWithExitCode
              "cc"
              ( ["-std=c11", "-Wall", "-Wextra", "-Werror"]
                  <> runtimeArguments
                  <> [snapshotRuntime, "-x", "c", "-", "-o", executable]
              )
              "int main(void) { return 0; }\n"
          assertEqual ("C compiler runtime diagnostics:\n" <> compilerErr) ExitSuccess compilerExit,
      testCase "semispace collector relocates a transitive global root" $
        withTempDirectory "aihc-arm64-semispace" $ \directory -> do
          runtimeArguments <- nativeRuntimeArguments RuntimeGcSemispace
          let executable = directory </> "semispace-check"
              source =
                unlines
                  [ "#include \"aihc_runtime.h\"",
                    "_Static_assert(sizeof(AihcValue) == sizeof(AihcSlot), \"one-word object header\");",
                    "static void entry_8(AihcSlot *arguments) { (void)arguments; }",
                    "static void entry_9(AihcSlot *arguments) { (void)arguments; }",
                    "static void entry_10(AihcSlot *arguments) { (void)arguments; }",
                    "static void entry_12(AihcSlot *arguments) { (void)arguments; }",
                    "static const uint8_t pointer_field[] = {1};",
                    "static const uint8_t pointer_then_scalar[] = {1, 0};",
                    "static const AihcInfo leaf_info = {1, 0, 0, 0, 0, 0, 0, AIHC_FRAME_NONE, AIHC_OBJECT_NODE};",
                    "static const AihcInfo box_info = {2, 0, 1, 0, pointer_field, 0, 0, AIHC_FRAME_NONE, AIHC_OBJECT_NODE};",
                    "static const AihcInfo thunk_info = {4, 0, 0, 0, 0, 0, 0, AIHC_FRAME_NONE, AIHC_OBJECT_THUNK};",
                    "static const AihcInfo partial_final_info = {3, 0, 2, 0, pointer_then_scalar, 0, 0, AIHC_FRAME_NONE, AIHC_OBJECT_NODE};",
                    "static const AihcInfo partial_one_info = {3, 0, 1, 1, pointer_then_scalar, &partial_final_info, 0, AIHC_FRAME_NONE, AIHC_OBJECT_PARTIAL_CONSTRUCTOR};",
                    "static const AihcInfo partial_info = {3, 0, 0, 2, 0, &partial_one_info, 0, AIHC_FRAME_NONE, AIHC_OBJECT_PARTIAL_CONSTRUCTOR};",
                    "static const AihcInfo continuation_final_info = {8, entry_8, 1, 0, pointer_field, 0, 0, AIHC_FRAME_NORMAL, AIHC_OBJECT_CLOSURE};",
                    "static const AihcInfo continuation_info = {8, entry_8, 0, 1, 0, &continuation_final_info, 0, AIHC_FRAME_NORMAL, AIHC_OBJECT_CLOSURE};",
                    "static const AihcInfo action_final_info = {9, entry_9, 0, 0, 0, 0, 0, AIHC_FRAME_NONE, AIHC_OBJECT_CLOSURE};",
                    "static const AihcInfo action_info = {9, entry_9, 0, 1, 0, &action_final_info, 0, AIHC_FRAME_NONE, AIHC_OBJECT_CLOSURE};",
                    "static const AihcInfo thread_done_final_info = {10, entry_10, 1, 0, pointer_field, 0, 0, AIHC_FRAME_STOP, AIHC_OBJECT_CLOSURE};",
                    "static const AihcInfo thread_done_info = {10, entry_10, 0, 1, 0, &thread_done_final_info, 0, AIHC_FRAME_STOP, AIHC_OBJECT_CLOSURE};",
                    "static const AihcInfo yield_final_info = {12, entry_12, 0, 0, 0, 0, 0, AIHC_FRAME_NONE, AIHC_OBJECT_CLOSURE};",
                    "static const AihcInfo yield_info = {12, entry_12, 0, 1, 0, &yield_final_info, 0, AIHC_FRAME_NONE, AIHC_OBJECT_CLOSURE};",
                    "static int test_apply_roots(void) {",
                    "  AihcMachine *machine = aihc_machine_new(0);",
                    "  AihcValue *leaf = aihc_make_node(machine, &leaf_info);",
                    "  AihcValue *partial = aihc_make_node(machine, &partial_info);",
                    "  AihcValue *continuation = aihc_make_node(machine, &continuation_info);",
                    "  (void)aihc_make_node(machine, &leaf_info);",
                    "  (void)aihc_make_node(machine, &leaf_info);",
                    "  (void)aihc_make_node(machine, &leaf_info);",
                    "  (void)aihc_make_node(machine, &leaf_info);",
                    "  AihcSlot argument = (AihcSlot)leaf;",
                    "  AihcValue *result = aihc_apply_slow(machine, partial, 1, &argument, &continuation);",
                    "  return aihc_value_kind(result) == AIHC_OBJECT_PARTIAL_CONSTRUCTOR &&",
                    "         aihc_value_arity(result) == 1 && aihc_value_count(result) == 1 &&",
                    "         aihc_value_info_table(result) == &partial_one_info &&",
                    "         aihc_value_info((AihcValue *)result->fields[0]) == 1 &&",
                    "         aihc_value_info(continuation) == 8;",
                    "}",
                    "static int test_scheduler_roots(void) {",
                    "  AihcMachine *machine = aihc_machine_new(1);",
                    "  AihcValue *thread_done = aihc_make_node(machine, &thread_done_info);",
                    "  AihcValue *action = aihc_make_node(machine, &action_info);",
                    "  AihcValue *yield = aihc_make_node(machine, &yield_info);",
                    "  aihc_set_thread_done_continuation(machine, thread_done);",
                    "  machine->globals[0] = (AihcSlot)yield;",
                    "  if (aihc_fork(machine, action) == 0) return 0;",
                    "  for (int index = 0; index < 5; ++index) {",
                    "    (void)aihc_make_node(machine, &leaf_info);",
                    "  }",
                    "  yield = (AihcValue *)machine->globals[0];",
                    "  const AihcResume *resume = aihc_yield(machine, yield);",
                    "  if (resume->kind != AIHC_RESUME_APPLY) return 0;",
                    "  thread_done = machine->thread_done_continuation;",
                    "  return aihc_value_info(resume->function) == 9 &&",
                    "         resume->continuation == thread_done &&",
                    "         aihc_value_info(thread_done) == 10;",
                    "}",
                    "static int test_blackhole_roots(void) {",
                    "  AihcMachine *machine = aihc_machine_new(2);",
                    "  AihcValue *thunk = aihc_make_node(machine, &thunk_info);",
                    "  AihcValue *leaf = aihc_make_node(machine, &leaf_info);",
                    "  machine->globals[0] = (AihcSlot)thunk;",
                    "  machine->globals[1] = (AihcSlot)leaf;",
                    "  aihc_begin_blackhole(machine, thunk);",
                    "  for (int index = 0; index < 8; ++index) {",
                    "    (void)aihc_make_node(machine, &leaf_info);",
                    "  }",
                    "  thunk = (AihcValue *)machine->globals[0];",
                    "  leaf = (AihcValue *)machine->globals[1];",
                    "  if (aihc_value_kind(thunk) != AIHC_OBJECT_BLACKHOLE) return 0;",
                    "  aihc_update_blackhole(machine, thunk, leaf);",
                    "  return aihc_value_kind(thunk) == AIHC_OBJECT_INDIRECTION &&",
                    "         thunk->fields[0] == (AihcSlot)leaf;",
                    "}",
                    "static int test_array_roots(void) {",
                    "  AihcMachine *machine = aihc_machine_new(1);",
                    "  AihcValue *leaf = aihc_make_node(machine, &leaf_info);",
                    "  AihcSlot root = (AihcSlot)leaf;",
                    "  aihc_ensure_heap(machine, 4, 1, &root);",
                    "  leaf = (AihcValue *)root;",
                    "  AihcValue *array = aihc_array_new(machine, 2, (AihcSlot)leaf);",
                    "  if (aihc_array_write(array, 1, (AihcSlot)leaf) != 0) return 0;",
                    "  machine->globals[0] = (AihcSlot)array;",
                    "  for (int index = 0; index < 100; ++index) {",
                    "    (void)aihc_make_node(machine, &leaf_info);",
                    "  }",
                    "  array = (AihcValue *)machine->globals[0];",
                    "  AihcValue *first = (AihcValue *)aihc_array_index(array, 0);",
                    "  AihcValue *second = (AihcValue *)aihc_array_index(array, 1);",
                    "  return aihc_value_kind(array) == AIHC_OBJECT_ARRAY &&",
                    "         aihc_array_same(array, array) == 1 &&",
                    "         first == second && aihc_value_info(first) == 1;",
                    "}",
                    "static int test_mutvar_roots(void) {",
                    "  AihcMachine *machine = aihc_machine_new(2);",
                    "  AihcValue *leaf = aihc_make_node(machine, &leaf_info);",
                    "  AihcSlot root = (AihcSlot)leaf;",
                    "  aihc_ensure_heap(machine, 3, 1, &root);",
                    "  leaf = (AihcValue *)root;",
                    "  AihcValue *first = aihc_mutvar_new(machine, (AihcSlot)leaf);",
                    "  machine->globals[0] = (AihcSlot)first;",
                    "  root = (AihcSlot)leaf;",
                    "  aihc_ensure_heap(machine, 3, 1, &root);",
                    "  leaf = (AihcValue *)root;",
                    "  AihcValue *second = aihc_mutvar_new(machine, (AihcSlot)leaf);",
                    "  machine->globals[1] = (AihcSlot)second;",
                    "  for (int index = 0; index < 100; ++index) {",
                    "    (void)aihc_make_node(machine, &leaf_info);",
                    "  }",
                    "  first = (AihcValue *)machine->globals[0];",
                    "  second = (AihcValue *)machine->globals[1];",
                    "  if (aihc_mutvar_write(first, aihc_mutvar_read(second)) != 0) return 0;",
                    "  leaf = (AihcValue *)aihc_mutvar_read(first);",
                    "  if (aihc_mutvar_compare_and_swap(first, (AihcSlot)leaf, (AihcSlot)second) != 0) return 0;",
                    "  if (aihc_mutvar_compare_and_swap(first, (AihcSlot)leaf, (AihcSlot)first) != 1) return 0;",
                    "  return aihc_mutvar_same(first, first) == 1 &&",
                    "         aihc_mutvar_same(first, second) == 0 &&",
                    "         aihc_mutvar_read(first) == (AihcSlot)second &&",
                    "         aihc_value_info(leaf) == 1;",
                    "}",
                    "int main(void) {",
                    "  if (!test_apply_roots()) return 1;",
                    "  if (!test_scheduler_roots()) return 1;",
                    "  if (!test_blackhole_roots()) return 1;",
                    "  if (!test_array_roots()) return 1;",
                    "  if (!test_mutvar_roots()) return 1;",
                    "  AihcMachine *machine = aihc_machine_new(1);",
                    "  AihcValue *leaf = aihc_make_node(machine, &leaf_info);",
                    "  if (leaf->header != (AihcSlot)(uintptr_t)&leaf_info) return 1;",
                    "  AihcValue *box = aihc_make_node(machine, &box_info);",
                    "  aihc_set_field(box, 0, (AihcSlot)leaf);",
                    "  machine->globals[0] = (AihcSlot)box;",
                    "  for (int index = 0; index < 100; ++index) {",
                    "    (void)aihc_make_node(machine, &leaf_info);",
                    "  }",
                    "  box = (AihcValue *)machine->globals[0];",
                    "  leaf = (AihcValue *)box->fields[0];",
                    "  return aihc_value_info(box) == 2 && aihc_value_info(leaf) == 1 ? 0 : 1;",
                    "}"
                  ]
          (compilerExit, _compilerOut, compilerErr) <-
            readProcessWithExitCode
              "cc"
              ( [ "-std=c11",
                  "-Wall",
                  "-Wextra",
                  "-Werror",
                  "-DAIHC_SEMISPACE_BYTES=64"
                ]
                  <> runtimeArguments
                  <> ["-x", "c", "-", "-o", executable]
              )
              source
          assertEqual ("C compiler semispace diagnostics:\n" <> compilerErr) ExitSuccess compilerExit
          (programExit, _programOut, programErr) <- readProcessWithExitCode executable [] ""
          assertEqual ("semispace runtime diagnostics:\n" <> programErr) ExitSuccess programExit,
      testCase "compiles standalone HelloWorld GRIN to native ARM64" testNativeHelloWorld,
      testCase "runs fork# and yield# with FIFO scheduling" testNativeScheduler,
      testCase "catches a synchronous exception" testNativeSynchronousException,
      testCase "blocks and wakes threads that enter a shared blackhole" testNativeBlackholeScheduler,
      testCase "waits for stdin and resumes an async stdio continuation" testNativeStdioScheduler
    ]

data SnapshotCase = SnapshotCase
  { snapshotCaseName :: !String,
    snapshotCaseProgram :: !GrinProgram,
    snapshotCaseEntry :: !FunctionName,
    snapshotCaseExpectation :: !SnapshotExpectation
  }

data SnapshotExpectation
  = SnapshotSuccess !T.Text !Word64
  | SnapshotFailure !T.Text

data SnapshotFixture = SnapshotFixture
  { snapshotFixtureEntry :: !T.Text,
    snapshotFixtureProgram :: !T.Text,
    snapshotFixtureReturn :: !(Maybe T.Text),
    snapshotFixtureHeap :: !(Maybe T.Text),
    snapshotFixtureError :: !(Maybe T.Text),
    snapshotFixtureAllocations :: !(Maybe (Map.Map T.Text Word64)),
    snapshotFixtureStatus :: !T.Text,
    snapshotFixtureReason :: !T.Text
  }

instance FromJSON SnapshotFixture where
  parseJSON =
    withObject "GRIN snapshot fixture" $ \object ->
      SnapshotFixture
        <$> object .: "entry"
        <*> object .: "program"
        <*> object .:? "return"
        <*> object .:? "heap"
        <*> object .:? "error"
        <*> object .:? "allocations"
        <*> object .: "status"
        <*> object .: "reason"

snapshotCases :: [(String, FilePath)]
snapshotCases =
  [ ("stores one value", "store-one.yaml"),
    ("preserves a suspended thunk", "store-suspended.yaml"),
    ("stores linked values", "store-linked.yaml"),
    ("stores a self-referential value", "store-self-referential.yaml"),
    ("returns an unboxed value", "return-unboxed.yaml"),
    ("loops ten million times", "loop-add.yaml"),
    ("evaluates only through GrinEval", "eval.yaml"),
    ("preserves WHNF pointers through GrinEval", "eval-whnf.yaml"),
    ("rejects blackholed thunk re-entry", "eval-blackhole.yaml"),
    ("saturates a partial constructor through C", "apply-constructor.yaml"),
    ("allocates a multi-stage partial closure", "apply-multistage.yaml"),
    ("saturates a partial closure through registers", "apply-partial.yaml"),
    ("passes excess saturated arguments through the native stack", "apply-register-overflow.yaml"),
    ("passes direct-call arguments through registers and the native stack", "call-register-overflow.yaml"),
    ("applies a stored closure", "apply.yaml"),
    ("snapshots the ThreadId# returned by fork#", "fork.yaml"),
    ("snapshots child evaluation after yield#", "yield.yaml")
  ]

snapshotTest :: (String, FilePath) -> TestTree
snapshotTest (name, fixtureName) =
  testCase name $ do
    snapshotCase <- loadSnapshotCase name fixtureName
    let program = snapshotCaseProgram snapshotCase
        entry = snapshotCaseEntry snapshotCase
        expectation = snapshotCaseExpectation snapshotCase
    assertEqual "direct GRIN lint" [] (lintProgram program)
    interpreted <- interpretProgramFunctionSnapshot entry program
    assertInterpretedExpectation expectation interpreted
    let gc = expectGcGrin program
    assertEqual "GC-GRIN lint" [] (lintProgram (gcGrinProgram gc))
    observed <-
      case compileObservedFunction entry gc of
        Left err -> assertFailure ("native snapshot compilation failed: " <> show err)
        Right value -> pure value
    when (fixtureName == "apply-partial.yaml") $ do
      let assembly = observedAssembly observed
      assertBool "dispatches through the info-table apply entry" ("ldr x8, [x8, #48]" `T.isInfixOf` assembly)
      assertBool
        "loads captured and supplied arguments into registers"
        ("mov x1, x0\n  ldr x0, [x20, #8]\n  b _aihc_snapshot_function_0" `T.isInfixOf` assembly)
    when (fixtureName == "apply-register-overflow.yaml") $ do
      let assembly = observedAssembly observed
      assertBool "spills supplied register overflow" ("sub sp, sp, x8\n  mov x9, sp" `T.isInfixOf` assembly)
      assertBool "reloads supplied register overflow" ("mov x9, sp\n  ldr x8, [x9], #8" `T.isInfixOf` assembly)
    when (fixtureName == "call-register-overflow.yaml") $ do
      let assembly = observedAssembly observed
      assertBool "spills direct-call register overflow" ("sub sp, sp, x8\n  mov x10, sp" `T.isInfixOf` assembly)
      assertBool "uses canonical direct-call entries" (not ("_register" `T.isInfixOf` assembly))
    when (fixtureName == "loop-add.yaml") $ do
      let assembly = observedAssembly observed
          loopAndRest = snd (T.breakOn "_aihc_snapshot_function_0:" assembly)
          loopAssembly = fst (T.breakOn "_aihc_snapshot_function_1:" loopAndRest)
      assertBool "keeps the loop's GRIN variables out of local spill storage" (not ("[x19" `T.isInfixOf` loopAssembly))
      assertEqual "only the self-tail-call branches to the allocated body" 1 (T.count "b _aihc_snapshot_function_0_body" loopAssembly)
      assertBool "merges case dispatch into the loop body" (not ("case_dispatch" `T.isInfixOf` loopAssembly))
      assertBool "uses the canonical label as the register entry" (not ("_register" `T.isInfixOf` loopAssembly))
    when (arch == "aarch64" && os == "darwin") $ do
      native <- runObservedProgram observed
      assertNativeExpectation expectation native

assertInterpretedExpectation :: SnapshotExpectation -> Either InterpretError HeapSnapshot -> IO ()
assertInterpretedExpectation expectation interpreted =
  case (expectation, interpreted) of
    (SnapshotSuccess expected _, Right snapshot) ->
      assertEqual "interpreter snapshot" (T.stripEnd expected) (T.stripEnd (renderHeapSnapshot snapshot))
    (SnapshotFailure expected, Left err) ->
      assertEqual "interpreter error" expected (renderInterpretFailure err)
    (SnapshotSuccess _ _, Left err) ->
      assertFailure ("GRIN interpreter failed: " <> show err)
    (SnapshotFailure _, Right snapshot) ->
      assertFailure ("GRIN interpreter unexpectedly succeeded:\n" <> T.unpack (renderHeapSnapshot snapshot))

assertNativeExpectation :: SnapshotExpectation -> Either T.Text T.Text -> IO ()
assertNativeExpectation expectation native =
  case (expectation, native) of
    (SnapshotSuccess expected expectedAllocations, Right snapshot) ->
      assertEqual
        "native snapshot"
        (T.stripEnd expected <> "\nallocations: " <> T.pack (show expectedAllocations))
        (T.stripEnd snapshot)
    (SnapshotFailure expected, Left err) ->
      assertEqual "native error" expected err
    (SnapshotSuccess _ _, Left err) ->
      assertFailure ("native snapshot failed: " <> T.unpack err)
    (SnapshotFailure _, Right snapshot) ->
      assertFailure ("native snapshot unexpectedly succeeded:\n" <> T.unpack snapshot)

renderInterpretFailure :: InterpretError -> T.Text
renderInterpretFailure err =
  case err of
    InterpretBlackhole _ -> "blackholed thunk re-entered"
    _ -> T.pack (show err)

loadSnapshotCase :: String -> FilePath -> IO SnapshotCase
loadSnapshotCase name fixtureName = do
  root <- snapshotFixtureRoot
  result <- Y.decodeFileEither (root </> fixtureName)
  fixture <-
    case result of
      Left err -> assertFailure ("invalid GRIN snapshot fixture: " <> Y.prettyPrintParseException err)
      Right value -> pure value
  assertEqual "fixture status" "pass" (snapshotFixtureStatus fixture)
  assertBool "fixture reason is present" (not (T.null (T.strip (snapshotFixtureReason fixture))))
  program <-
    case parseProgram (snapshotFixtureProgram fixture) of
      Left err -> assertFailure ("invalid GRIN program: " <> renderParseError err)
      Right value -> pure value
  expectation <-
    case (snapshotFixtureReturn fixture, snapshotFixtureHeap fixture, snapshotFixtureError fixture) of
      (Just returnValue, Just heapValue, Nothing) -> do
        let heap = T.stripEnd heapValue
            expected
              | heap == "[]" = "return: " <> returnValue <> "\nheap: []"
              | otherwise =
                  "return: "
                    <> returnValue
                    <> "\nheap:\n"
                    <> T.unlines (map ("  " <>) (T.lines heap))
        allocations <-
          case snapshotFixtureAllocations fixture >>= Map.lookup "macos-arm64" of
            Just count -> pure count
            Nothing -> assertFailure "successful snapshot fixture must define allocations.macos-arm64"
        pure (SnapshotSuccess expected allocations)
      (Nothing, Nothing, Just err)
        | not (T.null (T.strip err)) -> do
            assertEqual "failing snapshot allocations" Nothing (snapshotFixtureAllocations fixture)
            pure (SnapshotFailure (T.strip err))
      _ -> assertFailure "snapshot fixture must define either return and heap, or a non-empty error"
  pure
    SnapshotCase
      { snapshotCaseName = name,
        snapshotCaseProgram = program,
        snapshotCaseEntry = FunctionName (snapshotFixtureEntry fixture),
        snapshotCaseExpectation = expectation
      }

snapshotFixtureRoot :: IO FilePath
snapshotFixtureRoot = getCurrentDirectory >>= findRoot
  where
    findRoot directory = do
      let candidate = directory </> "test" </> "Test" </> "Fixtures" </> "grin-snapshot"
      exists <- doesDirectoryExist candidate
      if exists
        then pure candidate
        else do
          let parent = takeDirectory directory
          if parent == directory
            then assertFailure "GRIN snapshot fixture root is missing"
            else findRoot parent

runObservedProgram :: ObservedProgram -> IO (Either T.Text T.Text)
runObservedProgram observed =
  withTempDirectory "aihc-arm64-snapshot" $ \directory -> do
    runtimeArguments <- nativeRuntimeArguments RuntimeGcCalloc
    snapshotRuntime <- snapshotSourcePath
    let assemblyPath = directory </> "snapshot.s"
        metadataPath = directory </> "snapshot_metadata.c"
        executablePath = directory </> "snapshot"
    TIO.writeFile assemblyPath (observedAssembly observed)
    TIO.writeFile metadataPath (observedMetadataSource observed)
    (clangExit, _clangOut, clangErr) <-
      readProcessWithExitCode
        "clang"
        ( ["--target=" <> targetTriple, "-std=c11", "-Wall", "-Wextra", "-Werror"]
            <> runtimeArguments
            <> [snapshotRuntime, metadataPath, assemblyPath, "-o", executablePath]
        )
        ""
    case clangExit of
      ExitSuccess -> pure ()
      ExitFailure _ -> assertFailure ("clang failed to assemble observed GRIN:\n" <> clangErr)
    (programExit, programOut, programErr) <- readProcessWithExitCode executablePath [] ""
    case programExit of
      ExitSuccess -> do
        assertEqual "native stderr" "" programErr
        pure (Right (T.pack programOut))
      ExitFailure _ -> do
        assertEqual "native stdout" "" programOut
        pure (Left (renderNativeFailure (T.pack programErr)))

renderNativeFailure :: T.Text -> T.Text
renderNativeFailure stderr =
  let message = T.strip stderr
   in fromMaybe message (T.stripPrefix "aihc runtime: " message)

explicitEvaluationProgram :: GrinProgram
explicitEvaluationProgram =
  GrinProgram
    { grinConstructors = [],
      grinPrimitives = [],
      grinForeignCalls = [],
      grinExternalGlobals = [],
      grinExternalFunctions = [],
      grinWhnfGlobals = [],
      grinCafs = [],
      grinFunctions =
        [ GrinFunction
            { grinFunctionName = FunctionName "case_operand",
              grinFunctionLinkName = Nothing,
              grinFunctionParameters = [caseOperand],
              grinFunctionResultRep = BoxedRep Lifted,
              grinFunctionBody =
                GrinCase
                  (GrinVarValue caseOperand)
                  caseBinder
                  [GrinAlt GrinDefaultAlt [] (GrinConstant [GrinVarValue caseBinder])]
            },
          GrinFunction
            { grinFunctionName = FunctionName "apply_operand",
              grinFunctionLinkName = Nothing,
              grinFunctionParameters = [applyOperand],
              grinFunctionResultRep = BoxedRep Lifted,
              grinFunctionBody = GrinApply (BoxedRep Lifted) (GrinVarValue applyOperand) []
            }
        ]
    }
  where
    caseOperand = GrinVar "case_operand" 201 (BoxedRep Lifted)
    caseBinder = GrinVar "case_binder" 202 (BoxedRep Lifted)
    applyOperand = GrinVar "apply_operand" 203 (BoxedRep Lifted)

testNativeHelloWorld :: IO ()
testNativeHelloWorld = do
  cases <- loadEvalCases
  evalCase <-
    case find ((== "native-hello-world.yaml") . evalCaseId) cases of
      Just value -> pure value
      Nothing -> assertFailure "native HelloWorld fixture is missing"
  compileResult <- compileEvalCase evalCase
  fcProgram <-
    case compileResult of
      Right value -> pure value
      Left err -> assertFailure ("HelloWorld failed before GRIN lowering: " <> err)
  let grinProgram = lowerProgram fcProgram
  assertBool "GRIN has no unboxed-tuple nodes" (not ("(#,#)" `isInfixOf` renderProgram grinProgram))
  assembly <-
    case compileProgram (evalBindingNameInProgram fcProgram) (expectGcGrin grinProgram) of
      Right value -> pure value
      Left err -> assertFailure ("ARM64 lowering failed: " <> show err)
  let rendered = T.unpack assembly
  assertBool "emits indirect Haskell tail transfers" ("br x9" `isInfixOf` rendered)
  assertBool "never calls a generated Haskell entry" (not ("bl .Laihc_function_" `isInfixOf` rendered))
  assertBool "emits unboxed literals as raw words" (not ("_aihc_make_literal" `isInfixOf` rendered))
  when (arch == "aarch64" && os == "darwin") $
    runHelloWorldAssembly assembly

testNativeScheduler :: IO ()
testNativeScheduler = do
  assertEqual "direct GRIN lint" [] (lintProgram schedulerProgram)
  let gc = expectGcGrin schedulerProgram
  assertEqual "GC-GRIN lint" [] (lintProgram (gcGrinProgram gc))
  assembly <-
    case compileProgram "main" gc of
      Right value -> pure value
      Left err -> assertFailure ("ARM64 scheduler lowering failed: " <> show err)
  assertBool "captures argc and argv before machine startup" ("bl _aihc_program_arguments_initialize" `T.isInfixOf` assembly)
  assertBool "emits fork state operation" ("bl _aihc_fork" `T.isInfixOf` assembly)
  assertBool "emits yield state operation" ("bl _aihc_yield" `T.isInfixOf` assembly)
  assertBool "emits child completion transfer" ("bl _aihc_thread_done" `T.isInfixOf` assembly)
  let updateInfo = T.unlines (take 9 (dropWhile (/= ".Laihc_update_info:") (T.lines assembly)))
      finalInfo = T.unlines (take 9 (dropWhile (/= ".Laihc_final_info:") (T.lines assembly)))
  assertBool "emits update continuation frame metadata" ("  .quad 3" `T.isSuffixOf` T.stripEnd updateInfo)
  assertBool "emits stop continuation frame metadata" ("  .quad 5" `T.isSuffixOf` T.stripEnd finalInfo)
  when (arch == "aarch64" && os == "darwin") $
    runSchedulerAssembly "PCAB" assembly

testNativeSynchronousException :: IO ()
testNativeSynchronousException = do
  let gc = expectGcGrin synchronousExceptionProgram
  assembly <-
    case compileProgram "main" gc of
      Right value -> pure value
      Left err -> assertFailure ("ARM64 exception lowering failed: " <> show err)
  assertBool "emits the shared raise transfer" ("bl _aihc_raise" `T.isInfixOf` assembly)
  when (arch == "aarch64" && os == "darwin") $
    runSchedulerAssembly "E" assembly

testNativeBlackholeScheduler :: IO ()
testNativeBlackholeScheduler = do
  assertEqual "direct GRIN lint" [] (lintProgram blackholeSchedulerProgram)
  let gc = expectGcGrin blackholeSchedulerProgram
  assertEqual "GC-GRIN lint" [] (lintProgram (gcGrinProgram gc))
  assembly <-
    case compileProgram "main" gc of
      Right value -> pure value
      Left err -> assertFailure ("ARM64 blackhole scheduler lowering failed: " <> show err)
  when (arch == "aarch64" && os == "darwin") $
    runSchedulerAssembly "TA" assembly

testNativeStdioScheduler :: IO ()
testNativeStdioScheduler = do
  assertEqual "direct GRIN lint" [] (lintProgram stdioSchedulerProgram)
  let gc = expectGcGrin stdioSchedulerProgram
  assertEqual "GC-GRIN lint" [] (lintProgram (gcGrinProgram gc))
  assembly <-
    case compileProgram "main" gc of
      Right value -> pure value
      Left err -> assertFailure ("ARM64 stdio scheduler lowering failed: " <> show err)
  assertBool "emits generic IO suspension transfer" ("bl _aihc_await_io" `T.isInfixOf` assembly)
  assertBool "allocates a pinned byte array" ("bl _aihc_byte_array_new_pinned" `T.isInfixOf` assembly)
  assertBool "obtains the byte-array payload" ("bl _aihc_byte_array_contents" `T.isInfixOf` assembly)
  assertBool "obtains stdin through the runtime ABI" ("bl _aihc_io_stdin" `T.isInfixOf` assembly)
  assertBool "obtains stdout through the runtime ABI" ("bl _aihc_io_stdout" `T.isInfixOf` assembly)
  assertBool "submits a generic read through the runtime ABI" ("bl _aihc_io_submit_read" `T.isInfixOf` assembly)
  assertBool "submits a generic write through the runtime ABI" ("bl _aihc_io_submit_write" `T.isInfixOf` assembly)
  assertBool "consumes results through the runtime ABI" ("bl _aihc_io_take_result" `T.isInfixOf` assembly)
  assertBool "compiler has no operation-specific CPS transfer" (not ("_aihc_read_stdin_cps" `T.isInfixOf` assembly || "_aihc_write_stdout_cps" `T.isInfixOf` assembly))
  when (arch == "aarch64" && os == "darwin") $
    runStdioAssembly assembly

expectGcGrin :: GrinProgram -> GcGrinProgram
expectGcGrin program =
  case toCpsGrin program of
    Right cpsProgram -> lowerGc cpsProgram
    Left err -> error ("test GRIN failed CPS conversion: " <> show err)

runHelloWorldAssembly :: T.Text -> IO ()
runHelloWorldAssembly assembly =
  withTempDirectory "aihc-arm64-hello" $ \directory -> do
    runtimeArguments <- nativeRuntimeArguments RuntimeGcCalloc
    let assemblyPath = directory </> "hello.s"
        executablePath = directory </> "hello"
    writeFile assemblyPath (T.unpack assembly)
    (clangExit, _clangOut, clangErr) <-
      readProcessWithExitCode
        "clang"
        (["-std=c11", "-Wall", "-Wextra", "-Werror"] <> runtimeArguments <> [assemblyPath, "-o", executablePath])
        ""
    case clangExit of
      ExitSuccess -> pure ()
      ExitFailure _ -> assertFailure ("clang failed to assemble HelloWorld:\n" <> clangErr)
    (programExit, programOut, programErr) <- readProcessWithExitCode executablePath [] ""
    assertEqual ("native stderr: " <> programErr) ExitSuccess programExit
    assertEqual "native stdout" "Hello, world!\n" programOut

runSchedulerAssembly :: String -> T.Text -> IO ()
runSchedulerAssembly expected assembly =
  withTempDirectory "aihc-arm64-scheduler" $ \directory -> do
    runtimeArguments <- nativeRuntimeArguments RuntimeGcCalloc
    let assemblyPath = directory </> "scheduler.s"
        executablePath = directory </> "scheduler"
    TIO.writeFile assemblyPath assembly
    (clangExit, _clangOut, clangErr) <-
      readProcessWithExitCode
        "clang"
        (["-std=c11", "-Wall", "-Wextra", "-Werror"] <> runtimeArguments <> [assemblyPath, "-o", executablePath])
        ""
    assertEqual ("clang failed to assemble scheduler program:\n" <> clangErr) ExitSuccess clangExit
    (programExit, programOut, programErr) <- readProcessWithExitCode executablePath [] ""
    assertEqual ("native stderr: " <> programErr) ExitSuccess programExit
    assertEqual "scheduler stdout" expected programOut

runStdioAssembly :: T.Text -> IO ()
runStdioAssembly assembly =
  withTempDirectory "aihc-arm64-stdio" $ \directory -> do
    runtimeArguments <- nativeRuntimeArguments RuntimeGcCalloc
    let assemblyPath = directory </> "stdio.s"
        executablePath = directory </> "stdio"
    TIO.writeFile assemblyPath assembly
    (clangExit, _clangOut, clangErr) <-
      readProcessWithExitCode
        "clang"
        (["-std=c11", "-Wall", "-Wextra", "-Werror"] <> runtimeArguments <> [assemblyPath, "-o", executablePath])
        ""
    assertEqual ("clang failed to assemble async stdio program:\n" <> clangErr) ExitSuccess clangExit
    (Just childInput, Just childOutput, Just childError, processHandle) <-
      createProcess
        (proc executablePath [])
          { std_in = CreatePipe,
            std_out = CreatePipe,
            std_err = CreatePipe
          }
    threadDelay 50000
    hPutStr childInput "Buffered async IO\n"
    hFlush childInput
    hClose childInput
    programOut <- TIO.hGetContents childOutput
    programErr <- TIO.hGetContents childError
    programExit <- waitForProcess processHandle
    assertEqual ("native stderr: " <> T.unpack programErr) ExitSuccess programExit
    assertEqual "async stdout" "Buffered async IO\n" programOut

nativeRuntimeArguments :: RuntimeGarbageCollector -> IO [String]
nativeRuntimeArguments garbageCollector = do
  plan <- runtimePlan AppleArm64 garbageCollector
  pure
    ( ["-I" <> directory | directory <- runtimeIncludeDirectories plan]
        <> runtimeSources plan
    )

withTempDirectory :: String -> (FilePath -> IO value) -> IO value
withTempDirectory template = bracket acquire removeDirectoryRecursive
  where
    acquire = do
      temporary <- getTemporaryDirectory
      (path, handle) <- openTempFile temporary template
      hClose handle
      removeFile path
      createDirectory path
      pure path
