ReduceC.hs

{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE NoMonomorphismRestriction #-}

module ReduceC (
  defaultReduceC,
  defaultReduceCWithKeywords,
  -- reduceCTranslUnit,

  -- * Context
  Context (..),
  defaultContext,

  -- * Helpers
  prettyIdent,
) where

import Control.Applicative (Alternative (empty, (<|>)))
import Control.Monad (
  MonadPlus (mzero),
  foldM,
  forM,
  forM_,
  guard,
  join,
  mapAndUnzipM,
  unless,
  void,
  when,
 )
import qualified Control.Monad.IRTree as IRTree
import Control.Monad.Reduce (
  MonadReduce (split),
  collect,
  exceptIf,
  liftMaybe,
 )
import Control.Monad.State (
  MonadState (get, state),
  MonadTrans (lift),
  State,
  StateT (runStateT),
  evalState,
  evalStateT,
  gets,
  modify',
  runState,
 )
import Control.Monad.Trans.Maybe (MaybeT (runMaybeT))
import Data.Function ((&))
import Data.Functor (($>), (<&>))
import qualified Data.List as List
import qualified Data.List.NonEmpty as NonEmpty
import qualified Data.Map.Strict as Map
import Data.Maybe (
  catMaybes,
  fromMaybe,
  isJust,
  isNothing,
  mapMaybe,
 )
import qualified Data.Set as Set
import Data.Vector.Internal.Check (HasCallStack)
import qualified Language.C as C
import qualified Language.C.Data.Ident as C
import qualified Language.C.Data.Node as C

defaultReduceCWithKeywords :: (MonadReduce (String, C.Position) m) => [Keyword] -> C.CTranslUnit -> m C.CTranslUnit
defaultReduceCWithKeywords keywords a = reduceCTranslUnit a (defaultContext{keywords = Set.fromList keywords})
{-# SPECIALIZE defaultReduceCWithKeywords :: [Keyword] -> C.CTranslUnit -> IRTree.IRTree (String, C.Position) C.CTranslUnit #-}

defaultReduceC :: (MonadReduce (String, C.Position) m) => C.CTranslUnit -> m C.CTranslUnit
defaultReduceC a = reduceCTranslUnit a defaultContext
{-# SPECIALIZE defaultReduceC :: C.CTranslUnit -> IRTree.IRTree (String, C.Position) C.CTranslUnit #-}

reduceCTranslUnit ::
  (MonadReduce Lab m) =>
  C.CTranslationUnit C.NodeInfo ->
  Context ->
  m (C.CTranslationUnit C.NodeInfo)
reduceCTranslUnit (C.CTranslUnit es ni) ctx = do
  let _functions = foldMap (findFunctions (: [])) es

  let funmap :: [(C.Ident, Maybe Function)] =
        List.sortOn (maybe 0 (negate . funSize) . snd)
          . Map.toList
          . Map.fromListWith const
          . map (\f -> (funName f, Just f))
          . List.sortOn funSize
          $ _functions

  let reduce funcs = forM funcs \(k, mf) ->
        (k,) <$> runMaybeT do
          f <- liftMaybe mf
          let fstr = C.identToString (funName f)
          when (C.identToString (funName f) /= "main" || LoseMain `isIn` ctx) do
            exceptIf ("remove function " <> fstr <> " (" <> show (funSize f) <> ")", funPosition f)
          isStatic <-
            if funIsStatic f
              then
                split
                  ("remove static from " <> fstr, funPosition f)
                  (pure False)
                  (pure True)
              else pure False
          pure f{funIsStatic = isStatic}

  -- try remove static
  functions2 <- do
    funmap' <- reduce funmap
    if ComputeFunctionFixpoint `isIn` ctx
      then reduce funmap
      else pure funmap'

  functions3 <- forM functions2 \(k, mf) ->
    (k,) <$> runMaybeT do
      f <- liftMaybe mf
      if C.identToString (funName f) /= "main" || LoseMain `isIn` ctx
        then do
          params <- case funParams f of
            Just params -> do
              Just <$> forM (zip [1 :: Int ..] params) \(i, p) ->
                if p
                  then split ("remove parameter " <> show i <> " from " <> C.identToString (funName f), funPosition f) (pure False) (pure True)
                  else pure False
            ow -> pure ow
          pure f{funParams = params}
        else do
          pure f

  let builtins =
        [ ("fabsf", FunType (NonVoid TNum) (Params [Just TNum] False))
        , ("fabs", FunType (NonVoid TNum) (Params [Just TNum] False))
        , ("__builtin_abort", FunType Void (Params [] False))
        ]

  let functions''' =
        Map.fromList $
          [ ( funName
            , Just $
                Function
                  { funIsStatic = False
                  , funPosition = C.posOf funName
                  , funSize = 0
                  , funParams = case funTypeParams funType of
                      VoidParams -> Nothing
                      Params _ True -> Nothing
                      Params fx False -> Just [isJust f | f <- fx]
                  , ..
                  }
            )
          | (C.builtinIdent -> funName, funType) <- builtins
          ]
            <> functions3

  let ctx' =
        ctx
          { functions = functions'''
          , inlineExprs =
              inlineExprs ctx
                <> Map.fromList
                  [(C.builtinIdent f, IEKeep (TFun ft)) | (f, ft) <- builtins]
          }
  res' <- evalStateT (mapM reduceCExternalDeclaration es) ctx'
  pure $ C.CTranslUnit (catMaybes res') ni

newtype SpecifierFilter = SpecifierFilter
  { sfKeepStatic :: Bool
  }

keepAll :: SpecifierFilter
keepAll = SpecifierFilter{sfKeepStatic = True}

{- | Update the CDeclarationSpecifier's to match the context. Specifically, update
the typedefs and the structs. Alos return a base type.
-}
updateCDeclarationSpecifiers ::
  ( MonadState Context m
  , MonadPlus m
  ) =>
  SpecifierFilter ->
  [C.CDeclarationSpecifier C.NodeInfo] ->
  m (Voidable, [C.CDeclarationSpecifier C.NodeInfo])
updateCDeclarationSpecifiers sf spec = do
  ctx <- get
  spec' <- concat <$> mapM (updateSpec ctx) spec
  bt <- baseType ctx spec'
  pure (bt, spec')
 where