module Gyehoek.Test.Stack.VM (root) where

import Test.Tasty (TestTree, testGroup)
import Test.Tasty.HUnit
import Gyehoek.Stack.Syntax
import Gyehoek.Stack.VM qualified as Sut
import Data.List (List)


root :: IO TestTree
root = pure . testGroup "stack machine" $
  [ lit_int
  , procedure
  , prims
  ]



evalsTo :: List Obj -> List Block -> Assertion
evalsTo rs bs = Sut.eval (MkProgram bs) @?= rs



lit_int = testCase "lit int" do
  evalsTo [ObjImm (ImmInt 3)]
    [ MkBlock "main" []
      [ PopCont "ktail"
      , Call (ValReg "ktail") [ValImm (ImmInt 3)]
      ]
    ]

procedure = testGroup "procedure"
  [ testCase "return constant" do
      evalsTo [ObjImm (ImmInt 123)]
        [ MkBlock "main" []
          [ Call (ValLabel "silly") []
          ]
        , MkBlock "silly" []
          [ PopCont "ktail"
          , Call (ValReg "ktail") [ValImm (ImmInt 123)]
          ]
        ]
  , testCase "identity function" do
      evalsTo [ObjImm (ImmInt 45)]
        [ MkBlock "main" []
          [ Call (ValLabel "id") [ValImm (ImmInt 45)]
          ]
        , MkBlock "id" ["x"]
          [ PopCont "ktail"
          , Call (ValReg "ktail") [ValReg "x"]
          ]
        ]
  , testCase "square" do
      evalsTo [ObjImm (ImmInt 16)]
        [ MkBlock "main" []
          [ Call (ValLabel "square") [ValImm (ImmInt 4)]
          ]
        , MkBlock "square" ["x"]
          [ PopCont "ktail"
          , Prim "x2" $ PrimMul (ValReg "x") (ValReg "x")
          , Call (ValReg "ktail") [ValReg "x2"]
          ]
        ]
  , testCase "factorial" do
      let fac n =
            [ MkBlock "fac" ["n"]
              [ Prim "x0" $ PrimZeroP (ValReg "n")
              , If (ValReg "x0")
                [ PopCont "ktail"
                , Call (ValReg "ktail") [ValImm (ImmInt 1)]
                ]
                [ Push (ValReg "n")
                , Prim "x1" $ PrimSub (ValReg "n") (ValImm (ImmInt 1))
                , PushCont (ValLabel "fac-k0")
                , Call (ValLabel "fac") [ValReg "x1"]
                ]
              ]
            , MkBlock "fac-k0" ["x2"]
              [ Pop "n"
              , Prim "x3" $ PrimMul (ValReg "x2") (ValReg "n")
              , PopCont "ktail"
              , Call (ValReg "ktail") [ValReg "x3"]
              ]
            , MkBlock "main" []
              [ Call (ValLabel "fac") [ValImm (ImmInt n)]
              ]
            ]
      evalsTo [ObjImm (ImmInt 1)] $ fac 0
      evalsTo [ObjImm (ImmInt 720)] $ fac 6
  ]

prims = testGroup "prims"
  [ arith
  , testCase "zero?" do
      trivialPrimTest [ObjImm (ImmBool True)] $
        PrimZeroP $ ValImm $ ImmInt 0
      trivialPrimTest [ObjImm (ImmBool False)] $
        PrimZeroP $ ValImm $ ImmInt 12
  ]

trivialPrimTest rs p =
  evalsTo rs
  [ MkBlock "main" []
    [ PopCont "ktail"
    , Prim "x1" p
    , Call (ValReg "ktail") [ValReg "x1"]
    ]
  ]

arith = testGroup "arith"
  [ testCase "multipy" do
      trivialPrimTest [ObjImm (ImmInt 12)]
        (PrimMul (ValImm $ ImmInt 3) (ValImm $ ImmInt 4))
  , testCase "subtract" do
      trivialPrimTest [ObjImm (ImmInt 14)]
        (PrimSub (ValImm $ ImmInt 20) (ValImm $ ImmInt 6))
  ]