Main.hs

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

import Control.Monad.Reduce
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 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 <**> helper)
    . 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"
        , help "check for validity, throw error if command fails"
        ]

  pure
    $ usingLoggerT (cmap fmtMessage logTextStdout)
    $ do
      let
        test f = process D ("test " <> Text.pack f) do
          err <- liftIO $ runProcess (proc "clang" ["-O0", file])
          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 = do
          when validity do
            liftIO $ copyFile file (file <.> "last")
          output f c
          t <- test f
          if t
            then pure True
            else do
              liftIO $ when validity do
                copyFile file (file <.> "fail")
                copyFile (file <.> "last") file
                exitFailure
              pure False

      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 <-
        reduceI
          (check' file)
          (Map.singleton (C.internalIdent "main") True)
          (reduceC c)

      output file l
 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 -> do
    givenWith (funName fun)
    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 -> do
    givenWith x
    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 $ munder 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 -> do
    givenThat i
    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 -> onBothExpr 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 -> onBothExpr e1 e2 \e1' e2' ->
    pure $ C.CAssign op e1' e2' ni
  C.CIndex e1 e2 ni -> onBothExpr e1 e2 \e1' e2' ->
    pure $ C.CIndex e1' e2' ni
  C.CMember e i b ni -> do
    givenThat i
    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
  onBothExpr elhs erhs = onBoth (mrCExpression elhs) (mrCExpression erhs)

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)