Main.hs

{-# LANGUAGE ApplicativeDo #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}

import Control.RTree

import Colog
import Control.Applicative
import Control.Monad
import Control.Monad.Trans
import Control.Monad.Trans.Maybe
import Data.Foldable
import Data.Functor
import Data.Maybe
import Language.C qualified as C
import Options.Applicative
import System.Directory
import System.Exit
import System.FilePath

import Data.Text qualified as Text

-- import System.Process.Typed

import Data.Map.Strict qualified as Map
import GHC.Stack
import Language.C (CInitializer (CInitExpr))
import System.IO
import System.Process.Typed
import Text.Pretty.Simple
import Text.PrettyPrint qualified as P
import Prelude hiding (log)

main :: IO ()
main =
  join
    . execParser
    $ info run
    . fold
    $ []

process :: (WithLog env (Msg sev) m, MonadIO m) => sev -> Text.Text -> m a -> m a
process sev p ma = do
  msg "Started "
  a <- ma
  msg "Done "
  pure a
 where
  s = withFrozenCallStack callStack
  msg t = Colog.logMsg (Msg sev s (t <> p))

run :: (HasCallStack) => Parser (IO ())
run = do
  file <- strArgument $ fold [metavar "FILE"]

  validity <- flag False True $ fold [long "validity"]

  pure
    $ usingLoggerT (cmap fmtMessage logTextStdout)
    $ do
      let
        test f = process D ("test " <> Text.pack f) do
          err <- liftIO $ runProcess (proc "clang" ["-O0", "test.c"])
          log D (Text.pack $ show err)
          pure (err == ExitSuccess)

        output f c = process D ("writing " <> Text.pack f) do
          let x = P.render (C.pretty (c $> C.undefNode))
          liftIO $ writeFile f x

        check f _ c = MaybeT do
          when validity do
            liftIO $ copyFile file (file <.> "last")
          output f c
          t <- test f
          if t
            then pure (Just ())
            else do
              liftIO $ when validity do
                copyFile file (file <.> "fail")
                copyFile (file <.> "last") file
                exitFailure
              pure Nothing

      let bak = file <.> "bak"

      process D "backing up file" do
        liftIO $ copyFile file bak

      process I "check predicate" do
        t <- test file
        unless t do
          liftIO exitFailure

      c <- process D "parsing file" do
        parseCFile file

      output file c

      process I "sanity checks" do
        c' <- parseCFile file
        unless (void c' == void c) do
          liftIO do
            withFile "error.1.hs" WriteMode (`pHPrint` void c)
            withFile "error.2.hs" WriteMode (`pHPrint` void c')
          logError "Outputted a different file than i read... Please report original file and error.{1,2}.hs"
          liftIO exitFailure

        t <- test file
        unless t do
          liftIO exitFailure

      l <- runMaybeT (reduceIO (check file) (Map.singleton (C.internalIdent "main") True) (reduceC c))

      case l of
        Just l' -> do
          output file l'
          logInfo "Success"
        Nothing -> do
          logError "Unable to produce results"
          liftIO exitFailure
 where
  parseCFile file = do
    res <- liftIO $ C.parseCFilePre file
    case res of
      Right c -> pure c
      Left err -> do
        logError (Text.pack (show err))
        liftIO exitFailure

type Lab = C.Ident

reduceC :: (MonadReduce Lab m) => C.CTranslUnit -> m C.CTranslUnit
reduceC (C.CTranslUnit es ni) = do
  es' <- collect mrCExternalDeclaration es
  pure $ C.CTranslUnit es' ni

mrCExternalDeclaration :: (MonadReduce Lab m) => C.CExternalDeclaration C.NodeInfo -> MaybeT m (C.CExternalDeclaration C.NodeInfo)
mrCExternalDeclaration = \case
  C.CFDefExt fun ->
    split
      (funName fun)
      empty
      (C.CFDefExt <$> rCFunctionDef fun)
  C.CDeclExt decl ->
    C.CDeclExt <$> mrCDeclaration decl
  a -> error (show a)
 where
  funName (C.CFunDef _ (C.CDeclr x _ _ _ _) _ _ _) =
    x

mrCDeclarationItem :: (MonadReduce Lab m) => C.CDeclarationItem C.NodeInfo -> MaybeT m (C.CDeclarationItem C.NodeInfo)
mrCDeclarationItem = \case
  C.CDeclarationItem d@(C.CDeclr x _ _ _ _) i e ->
    split x empty do
      i' <- mtry $ munder i mrCInitializer
      e' <- mtry $ munder e mrCExpression
      pure (C.CDeclarationItem d i' e')
  a -> error (show a)

mrCInitializer :: (MonadReduce Lab m) => C.CInitializer C.NodeInfo -> MaybeT m (C.CInitializer C.NodeInfo)
mrCInitializer = \case
  C.CInitExpr e ni -> mrCExpression e <&> \e' -> CInitExpr e' ni
  C.CInitList (C.CInitializerList items) ni -> do
    collectNonEmpty' rmCInitializerListItem items <&> \items' ->
      C.CInitList (C.CInitializerList items') ni
 where
  rmCInitializerListItem (pds, is) = do
    pds' <- lift $ collect rmCPartDesignator pds
    is' <- mrCInitializer is
    pure (pds', is')

  rmCPartDesignator :: (MonadReduce Lab m) => C.CPartDesignator C.NodeInfo -> m (C.CPartDesignator C.NodeInfo)
  rmCPartDesignator = \case
    a -> error (show a)

mrCDeclaration :: (MonadReduce Lab m) => C.CDeclaration C.NodeInfo -> MaybeT m (C.CDeclaration C.NodeInfo)
mrCDeclaration = \case
  C.CDecl spc decl ni -> do
    decl' <- lift $ collect mrCDeclarationItem decl
    case decl' of
      [] -> empty
      decl'' -> pure $ C.CDecl spc decl'' ni
  a -> error (show a)

rCFunctionDef :: (MonadReduce Lab m) => C.CFunctionDef C.NodeInfo -> m (C.CFunctionDef C.NodeInfo)
rCFunctionDef (C.CFunDef spc dec cdecls smt ni) = do
  smt' <- rCStatement smt
  pure $ C.CFunDef spc dec cdecls smt' ni

rCStatement :: (MonadReduce Lab m) => C.CStatement C.NodeInfo -> m (C.CStatement C.NodeInfo)
rCStatement = \case
  C.CCompound is cbi ni -> do
    cbi' <- collect mrCCompoundBlockItem cbi
    pure $ C.CCompound is cbi' ni
  C.CExpr e ni -> do
    e' <- runMaybeT $ mlift e >>= mrCExpression
    pure $ C.CExpr e' ni
  C.CIf e s els ni -> do
    e' <- runMaybeT $ mrCExpression e
    s' <- rCStatement s
    els' <- case els of
      Just els' -> do
        pure Nothing <| Just <$> rCStatement els'
      Nothing -> pure Nothing
    case (e', els') of
      (Nothing, Nothing) -> pure s'
      (Just e'', Nothing) -> pure $ C.CIf e'' s' Nothing ni
      (Nothing, Just x) -> pure $ C.CIf zeroExp s' (Just x) ni
      (Just e'', Just x) -> pure $ C.CIf e'' s' (Just x) ni
  C.CFor e1 e2 e3 s ni -> do
    rCStatement s <| do
      e1' <- rCForInit e1
      e2' <- runMaybeT $ munder e2 mrCExpression
      e3' <- runMaybeT $ munder e3 mrCExpression
      s' <- rCStatement s
      pure $ C.CFor e1' e2' e3' s' ni
  C.CReturn e ni -> do
    e' <- case e of
      Nothing -> pure Nothing
      Just e' -> Just <$> zrCExpression e'
    pure $ C.CReturn e' ni
  C.CBreak ni -> pure (C.CBreak ni)
  C.CCont ni -> pure (C.CCont ni)
  C.CLabel i s [] ni ->
    -- todo fix attrs
    splitOn i (rCStatement s) do
      s' <- rCStatement s
      pure $ C.CLabel i s' [] ni
  C.CGoto i ni ->
    -- todo fix attrs
    splitOn i (pure $ C.CExpr Nothing ni) do
      pure $ C.CGoto i ni
  a -> error (show a)
 where
  rCForInit = \case
    C.CForDecl decl -> do
      m <- runMaybeT $ mrCDeclaration decl
      pure $ case m of
        Nothing -> C.CForInitializing Nothing
        Just d' -> C.CForDecl d'
    C.CForInitializing n -> do
      C.CForInitializing <$> runMaybeT (munder n mrCExpression)

orZero :: Maybe (C.CExpression C.NodeInfo) -> C.CExpression C.NodeInfo
orZero = fromMaybe zeroExp

zeroExp :: C.CExpression C.NodeInfo
zeroExp = C.CConst (C.CIntConst (C.cInteger 0) C.undefNode)

zrCExpression :: (MonadReduce Lab m) => C.CExpression C.NodeInfo -> m (C.CExpression C.NodeInfo)
zrCExpression e = orZero <$> runMaybeT (mrCExpression e)

mrCExpression :: (MonadReduce Lab m) => C.CExpression C.NodeInfo -> MaybeT m (C.CExpression C.NodeInfo)
mrCExpression = \case
  C.CVar i ni ->
    splitOn i empty (pure $ C.CVar i ni)
  C.CCall e es ni -> do
    e' <- mrCExpression e
    es' <- lift $ traverse zrCExpression es
    pure $ C.CCall e' es' ni
  C.CCond ec et ef ni -> do
    ec' <- mrCExpression ec
    ef' <- mrCExpression ef
    et' <- mtry $ munder et mrCExpression
    pure $ C.CCond ec' et' ef' ni
  C.CBinary o elhs erhs ni -> ejoin elhs erhs \lhs rhs ->
    pure $ C.CBinary o lhs rhs ni
  C.CUnary o elhs ni -> do
    lhs <- mrCExpression elhs
    pure $ C.CUnary o lhs ni
  C.CConst c -> do
    -- TODO fix
    pure $ C.CConst c
  C.CCast cd e ni -> do
    -- TODO fix
    cd' <- mrCDeclaration cd
    e' <- mrCExpression e
    pure $ C.CCast cd' e' ni
  C.CAssign op e1 e2 ni -> ejoin e1 e2 \e1' e2' ->
    pure $ C.CAssign op e1' e2' ni
  C.CIndex e1 e2 ni -> ejoin e1 e2 \e1' e2' ->
    pure $ C.CIndex e1' e2' ni
  C.CMember e i b ni -> do
    splitOn i empty do
      e' <- mrCExpression e
      pure $ C.CMember e' i b ni
  C.CComma items ni -> do
    C.CComma <$> collectNonEmpty' mrCExpression items <*> pure ni
  e -> error (show e)
 where
  ejoin elhs erhs = mjoin (mrCExpression elhs) (mrCExpression erhs)

mjoin :: (Monad m) => MaybeT m a -> MaybeT m a -> (a -> a -> MaybeT m a) -> MaybeT m a
mjoin mlhs mrhs fn = MaybeT do
  lhs <- runMaybeT mlhs
  case lhs of
    Nothing -> runMaybeT mrhs
    Just l -> do
      rhs <- runMaybeT mrhs
      case rhs of
        Nothing -> pure (Just l)
        Just r -> runMaybeT (fn l r)

mtry :: (Functor m) => MaybeT m a -> MaybeT m (Maybe a)
mtry (MaybeT mt) = MaybeT (Just <$> mt)

mlift :: (Applicative m) => Maybe a -> MaybeT m a
mlift a = MaybeT (pure a)

munder :: (Monad m) => Maybe a -> (a -> MaybeT m b) -> MaybeT m b
munder a mf = mlift a >>= mf

mrCCompoundBlockItem :: (MonadReduce Lab m) => C.CCompoundBlockItem C.NodeInfo -> MaybeT m (C.CCompoundBlockItem C.NodeInfo)
mrCCompoundBlockItem = \case
  C.CBlockStmt s -> empty <| lift (C.CBlockStmt <$> rCStatement s)
  C.CBlockDecl d -> C.CBlockDecl <$> mrCDeclaration d
  a -> error (show a)