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 ViewPatterns #-}
{-# LANGUAGE NoMonomorphismRestriction #-}

module ReduceC (
  defaultReduceC,
  defaultReduceCWithKeywords,
  -- reduceCTranslUnit,

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

  -- * Helpers
  prettyIdent,
) where

import Control.Applicative
import Control.Monad
import qualified Control.Monad.IRTree as IRTree
import Control.Monad.Reduce
import Control.Monad.State
import Control.Monad.Trans.Maybe
import Data.Data
import Data.Foldable
import Data.Function
import Data.Functor
import qualified Data.List as List
import qualified Data.Map.Strict as Map
import Data.Maybe
import Data.Monoid
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

reduceCTranslUnit
  :: (MonadReduce Lab m)
  => C.CTranslationUnit C.NodeInfo
  -> Context
  -> m (C.CTranslationUnit C.NodeInfo)
reduceCTranslUnit (C.CTranslUnit es ni) ctx = do
  (_functions, _structs) <- flip evalState ctx do
    (fs, sts) <- flip mapAndUnzipM es \e -> do
      includeTypeDef e
      funcs <- gets \ctx' -> findFunctions (: []) ctx' e
      structs <- gets \ctx' -> findStructs (: []) ctx' e
      pure (funcs, structs)
    pure (pure (concat fs, concat sts))

  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 <- reduce =<< reduce funmap

  functions3 <- forM functions2 \(k, mf) ->
    (k,) <$> runMaybeT do
      f <- liftMaybe mf
      params <- case funParams f of
        Params params False -> do
          params' <- forM params \p -> runMaybeT do
            p' <- liftMaybe p
            exceptIf ("remove parameter", funPosition f)
            pure p'
          pure (Params params' False)
        ow -> pure ow
      pure f{funParams = params}

  let functions''' =
        Map.fromList $
          functions3
            <> [ ( funName
                 , Just $
                    Function
                      { funIsStatic = False
                      , funPosition = C.posOf funName
                      , funSize = 0
                      , ..
                      }
                 )
               | (C.builtinIdent -> funName, funReturns, funParams) <-
                  [ ("fabsf", (Just CTNum), (Params [Just CTNum] False))
                  , ("fabs", (Just CTNum), (Params [Just CTNum] False))
                  ]
               ]

  structs' <- flip execStateT (structs ctx) do
    forM_ _structs \s ->
      modify' (Map.insert (structName s) (Just s))

  let ctx' = ctx{functions = functions''', structs = structs'}
  res' <- evalStateT (mapM reduceCExternalDeclaration es) ctx'
  pure $ C.CTranslUnit (catMaybes res') ni

reduceCExternalDeclaration
  :: (HasCallStack, MonadReduce Lab m)
  => C.CExternalDeclaration C.NodeInfo
  -> StateT Context m (Maybe (C.CExternalDeclaration C.NodeInfo))
reduceCExternalDeclaration r = case r of
  C.CFDefExt (C.CFunDef spec declr [] stmt ni) -> runMaybeT do
    ctx <- get
    guard (not $ any (shouldDeleteDeclSpec ctx) spec)

    let C.CDeclr mid dd Nothing [] ni2 = declr
    let (C.CFunDeclr (C.CFunParamsNew params b) attr ni3 : dd') = dd

    (pFilter, spec') <- case mid of
      Just fid -> do
        f <- liftMaybe (lookupFunction ctx fid)
        pure (funParams f, filterStorageModifiers (funIsStatic f) spec)
      Nothing -> do
        exceptIf ("remove function", C.posOf r)

        case params of
          [C.CDecl [C.CTypeSpec (C.CVoidType _)] [] _] -> pure (VoidParams, spec)
          _ow -> pure (Params (Just . snd <$> map (functionParameter ctx) params) False, spec)

    let (params', idents) = case pFilter of
          Params flt False -> filterParams ctx flt params
          _ow -> (params, [])

    labs <- flip collect (labelsOf stmt) \l -> do
      exceptIf ("remove label" <> show l, C.posOf l)
      pure l

    stmt' <-
      reduceCStatementOrEmptyBlock stmt StmtContext{stmtLabels = labs, stmtInLoop = False} $
        foldr (uncurry addInlineExpr) ctx idents

    let dd'' = C.CFunDeclr (C.CFunParamsNew params' b) attr ni3 : dd'

    pure . C.CFDefExt $
      C.CFunDef
        (inlineTypeDefsSpecs spec' ctx)
        (inlineTypeDefsCDeclarator (C.CDeclr mid dd'' Nothing [] ni2) ctx)
        []
        stmt'
        ni

  -- Type definitions
  C.CDeclExt d@(C.CDecl (C.CStorageSpec (C.CTypedef _) : rst) [item] ni) -> runMaybeT do
    let C.CDeclarationItem (C.CDeclr (Just ix) [] Nothing [] _) Nothing Nothing = item
    split
      ("Inline typedef" <> C.identToString ix, C.posOf ni)
      (modify (\ctx' -> addTypeDefs [ix] (ctype ctx' rst, ITInline rst) ctx') >> empty)
      do
        modify (\ctx' -> addTypeDefs [ix] (ctype ctx' rst, ITKeep) ctx')
        gets (C.CDeclExt <$> inlineTypeDefsCDeclaration d)

  -- The rest.
  C.CDeclExt (C.CDecl spec items ni) -> runMaybeT do
    ctx <- get

    let t = ctype ctx spec

    lift $ includeTypeDef r

    keep <- containsStructDeclaration spec

    -- Try to remove each declaration item
    (items', or -> isStatic) <-
      unzip <$> flip collect items \case
        di@(C.CDeclarationItem (C.CDeclr mid dd Nothing [] ni2) einit size) -> do
          case dd of
            (C.CFunDeclr params attr ni3) : rst -> do
              (dd', isStatic) <- case mid of
                Just fid -> do