------------------------------------------------------------------------------
-- | Author:  Michael Hanus
--   Version: November 2025
--
-- Some support to test the library `FlatCurry.AddTypes`.
--
------------------------------------------------------------------------------

module TestAddTypes
 where

import System.Environment ( getArgs )

import FlatCurry.AddTypes ( AExpr(..), ABranchExpr(..), TypeInfo, tiType
                          , annProg2Prog, addTypesInProg, fromAExpr)
import FlatCurry.Files    ( readFlatCurry, writeFlatCurryFile )
import FlatCurry.Pretty   ( showPrettyProg )
import FlatCurry.Types

------------------------------------------------------------------------------
-- For testing:

main :: IO ()
main = do
  args <- getArgs
  case args of
    [mname,out] -> do prog <- readFlatCurry mname
                      nprog <- decorateProgWithTypes prog
                      writeFlatCurryFile out nprog
                      putStrLn $ "Transformed fcy written into " ++ out
    [mname]     -> do prog <- readFlatCurry mname
                      nprog <- decorateProgWithTypes prog
                      putStrLn $ showPrettyProg nprog
    _ -> putStrLn $ "Wrong arguments! Use ...mod [fcyoutfile]"

test :: IO ()
test = testProg "Test"

testPrelude :: IO ()
testPrelude = testProg "Prelude"

testProg :: String -> IO ()
testProg mname = do
  prog <- readFlatCurry mname
  decorateProgWithTypes prog>>= putStrLn . showPrettyProg

-- | Add type information to the body of each function of the given program
--   and transform the annotated body to make the type information visible.
decorateProgWithTypes :: Prog -> IO Prog
decorateProgWithTypes prog =
  fmap (annProg2Prog (fromAExpr . addExplicitTypesInExp))
       (addTypesInProg prog)

-- Decorates an annotated expression with explicit `Typed` constructs
-- in order to make the type annotations visiable in the FlatCurry program.
-- This transformation is useful for testing.
addExplicitTypesInExp :: AExpr TypeInfo -> AExpr TypeInfo
addExplicitTypesInExp = addExp
 where
  addExp exp = case exp of
    AVar  ti _        -> ATyped ti exp (tiType ti)
    ALit  ti _        -> ATyped ti exp (tiType ti)
    AComb ti ct qn es -> ATyped ti (AComb ti ct qn (map addExp es)) (tiType ti)
    ALet  ti bs e     -> ALet ti (map (\ (v,t,be) -> (v, t, addExp be)) bs)
                                 (addExp e)
    AFree ti vs e     -> ATyped ti (AFree ti vs (addExp e)) (tiType ti)
    AOr   ti e1 e2    -> ATyped ti (AOr ti (addExp e1) (addExp e2)) (tiType ti)
    ACase ti ct ce bs -> ACase ti ct (addExp ce)
                          (map (\ (ABranch pt be) -> ABranch pt (addExp be)) bs)
    ATyped _ _ _      -> exp

-----------------------------------------------------------------------