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
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.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
import qualified Data.Set as Set
import Data.Vector.Internal.Check (HasCallStack)
import Debug.Pretty.Simple
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))
        ]

  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
  baseType
    :: (MonadPlus m)
    => Context
    -> [C.CDeclarationSpecifier C.NodeInfo]
    -> m Voidable
  baseType ctx =
    liftMaybe
      . exactlyOne
      . map \case
        C.CVoidType _ -> Just Void
        C.CSUType c _ -> NonVoid . TStruct <$> structId c
        C.CCharType _ -> Just $ NonVoid TNum
        C.CShortType _ -> Just $ NonVoid TNum
        C.CIntType _ -> Just $ NonVoid TNum
        C.CFloatType _ -> Just $ NonVoid TNum
        C.CDoubleType _ -> Just $ NonVoid TNum
        C.CSignedType _ -> Just $ NonVoid TNum
        C.CUnsigType _ -> Just $ NonVoid TNum
        C.CBoolType _ -> Just $ NonVoid TNum
        C.CLongType _ -> Just $ NonVoid TNum
        C.CInt128Type _ -> Just $ NonVoid TNum
        C.CFloatNType{} -> Just $ NonVoid TNum
        C.CEnumType (C.CEnum (Just ix) _ _ _) _ ->
          NonVoid TNum
            <$ guard (lookupEnum ctx ix == INKeep)
        C.CEnumType (C.CEnum Nothing _ _ _) _ -> Just $ NonVoid TNum
        C.CTypeDef idx _ ->
          case Map.lookup idx (typeDefs ctx) of
            Just (ITKeep t') -> Just t'
            Just ITDelete -> Nothing
            Just (ITInline t' _) -> Just t'
            Nothing -> error "error"
        a -> notSupportedYet (void a) a
      . typeSpecs
   where
    typeSpecs = mapMaybe \case
      C.CTypeSpec ts -> Just ts
      _ow -> Nothing

    exactlyOne =
      maybe
        (error "no type in type-specs")
        ( \case
            (t, []) -> NonEmpty.head t
            (t, rs) -> error ("more than one type in type-specs: " <> show (t : rs))
        )
        . List.uncons
        . NonEmpty.group

    structId (C.CStruct t mi md _ ni) =
      case mi of
        Just ix -> case lookupStruct ctx ix of
          ISDelete -> Nothing
          _ow -> Just $ Left ix
        Nothing ->
          let p' =
                maybe
                  (error $ "invalid struct at" <> show (C.posOf ni))
                  (concatMap namesAndTypeOf)
                  md
           in pure $ Right (StructType t Nothing p')

    -- structTypeOf (C.CStruct t mi md _ ni) =
    --   case mi of
    --     Just ix -> lookupStruct ctx ix
    --     Nothing ->
    --       let p' = maybe (error $ "invalid struct at" <> show (C.posOf ni)) (concatMap namesAndTypeOf) md
    --        in Just $ StructType t mi (Just p')

    namesAndTypeOf = \case
      C.CDecl spec2 items ni ->
        flip map items \case
          C.CDeclarationItem decl@(C.CDeclr (Just ix) _ _ _ _) _ _ ->
            (ix, nonVoid <$> typeOf spec2 decl)
          a -> notSupportedYet (void a) ni
      a -> notSupportedYet' a

    typeOf spec2 decl = baseType ctx spec2 >>= extendTypeWith decl

    extendTypeWith (C.CDeclr _ dd _ _ _) t =
      foldr applyDD (Just t) dd
     where
      applyDD = \case
        C.CPtrDeclr _ _ -> fmap (NonVoid . TPointer)
        C.CArrDeclr{} -> fmap (NonVoid . TPointer)
        C.CFunDeclr params _ ni -> \c ->
          case params of
            C.CFunParamsNew params' varadic -> do
              c' <- c
              Just $ NonVoid $ TFun (FunType c' (findParams varadic params'))
            b -> notSupportedYet b ni

      findParams varadic = \case
        [C.CDecl [C.CTypeSpec (C.CVoidType _)] [] _] -> VoidParams
        rst -> flip Params varadic $ flip map rst \case
          C.CDecl spec' [] _ ->
            nonVoid <$> baseType ctx spec'
          C.CDecl spec' [C.CDeclarationItem decl _ _] _ ->
            nonVoid <$> typeOf spec' decl
          a -> notSupportedYet' a

  updateSpec ctx a = case a of
    C.CTypeSpec t -> case t of
      C.CSUType (C.CStruct st (Just i) (Just declrs) attr x) b -> do
        fields <- case lookupStruct ctx i of
          ISDelete -> empty
          ISDeclared _ -> empty
          ISKeep s -> do
            pure $ structTypeFields s
        let declrs' :: [C.CDeclaration C.NodeInfo] = filterStruct ctx fields declrs
        pure [C.CTypeSpec (C.CSUType (C.CStruct st (Just i) (Just declrs') attr x) b)]
      C.CTypeDef idx _ -> do
        case Map.lookup idx . typeDefs $ ctx of
          Just (ITKeep _) -> pure [C.CTypeSpec t]
          Just (ITInline _ res) -> pure res
          Just ITDelete -> mzero
          Nothing -> error ("could not find typedef: " <> show idx)
      _ow -> pure [C.CTypeSpec t]
    C.CStorageSpec (C.CStatic _) -> pure [a | sfKeepStatic sf]
    C.CFunSpec (C.CInlineQual _) -> pure [a | sfKeepStatic sf]
    _ow -> pure [a]

  filterStruct ctx fields declrs =
    flip evalState fields do
      declrs' <- forM declrs $ \case
        C.CDecl spec2 items l -> runMaybeT do
          items' <- forM items $ \case
            C.CDeclarationItem (C.CDeclr mid dd sl attr ni2) enit ni1 -> runMaybeT do
              _ <- liftMaybe =<< state (\((_, t) : tps) -> (t, tps))
              (_, dd') <- liftMaybe (evalStateT (updateCDerivedDeclarators Void (repeat True) dd) ctx)
              pure (C.CDeclarationItem (C.CDeclr mid dd' sl attr ni2) enit ni1)
            a' -> notSupportedYet a' l
          (_, spec2') <- liftMaybe (evalStateT (updateCDeclarationSpecifiers keepAll spec2) ctx)
          let items'' = catMaybes items'
          guard $ not (List.null items'')
          pure (C.CDecl spec2' items'' l)
        a' -> notSupportedYet' a'
      pure $ catMaybes declrs'

updateCDerivedDeclarators
  :: forall m
   . ( MonadState Context m
     , MonadPlus m
     )
  => Voidable
  -> [Bool]
  -> [C.CDerivedDeclarator C.NodeInfo]
  -> m (Voidable, [C.CDerivedDeclarator C.NodeInfo])
updateCDerivedDeclarators bt ff dd = do
  foldM applyDD (bt, []) (reverse dd)
 where
  applyDD
    :: (r ~ (Voidable, [C.CDerivedDeclarator C.NodeInfo]))
    => r
    -> C.CDerivedDeclarator C.NodeInfo
    -> m r
  applyDD (t, dd') d = case d of
    C.CPtrDeclr _ _ -> do
      pure (NonVoid . TPointer $ t, d : dd')
    C.CArrDeclr{} ->
      pure (NonVoid . TPointer $ t, d : dd')
    C.CFunDeclr params arr ni -> do
      case params of
        C.CFunParamsNew params' varadic -> do
          (tp, params'') <- state (runState (findParams varadic params'))
          let t' = NonVoid $ TFun (FunType t tp)
          pure (t', C.CFunDeclr (C.CFunParamsNew params'' varadic) arr ni : dd')
        b -> notSupportedYet b ni

  findParams
    :: Bool
    -> [C.CDeclaration C.NodeInfo]
    -> State Context (Params, [C.CDeclaration C.NodeInfo])
  findParams varadic decls = case decls of
    [C.CDecl [C.CTypeSpec (C.CVoidType _)] [] _] ->
      pure (VoidParams, decls)
    _ow -> flip evalStateT ff do
      result <-
        forM decls $ \case
          C.CDecl spec items ni -> do
            keep <- state (\(t : tps) -> (t, tps))
            lift . runMaybeT $ do
              markDeleted items
              (bt', spec') <- updateCDeclarationSpecifiers keepAll spec
              (t, items') <- case items of
                [] -> do
                  guard keep
                  pure (nonVoid bt', [])
                [C.CDeclarationItem (C.CDeclr mid dd2 Nothing [] ni3) Nothing ni2] -> do
                  (t, dd2') <- case mid of
                    Just ix -> do
                      (nonVoid -> t, dd2') <- updateCDerivedDeclarators bt' [] dd2
                      guard keep
                      modify' (addInlineExpr ix (IEKeep t))
                      pure (t, dd2')
                    Nothing -> do
                      (nonVoid -> t, dd2') <- updateCDerivedDeclarators bt' [] dd2
                      guard keep
                      pure (t, dd2')
                  pure (t, [C.CDeclarationItem (C.CDeclr mid dd2' Nothing [] ni3) Nothing ni2])
                _ow -> notSupportedYet items ni
              pure (t, C.CDecl spec' items' ni)
          a -> notSupportedYet' a
      let (ts, decls') = unzip $ flip map result \case
            Just (t, d') -> (Just t, [d'])
            Nothing -> (Nothing, [])
      pure (Params ts varadic, concat decls')

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

    let C.CDeclr mid dd Nothing [] ni2 = declr

    mfun <- case mid of
      Just fid -> do
        modify' (addInlineExpr fid IEDelete)
        Just <$> liftMaybe (lookupFunction ctx fid)
      Nothing ->
        pure Nothing

    let keepStatic = maybe True funIsStatic mfun
    -- TODO handle this edgecase (struct declared in function declaration)
    _ <- reduceStructDeclaration spec
    (bt, spec') <- updateCDeclarationSpecifiers keepAll{sfKeepStatic = keepStatic} spec
    ((nonVoid -> t@(TFun (FunType rt _)), dd'), ctx') <-
      runStateT
        (updateCDerivedDeclarators bt (fromMaybe (repeat True) (mfun >>= funParams)) dd)
        ctx

    case mfun of
      Just fun -> do
        modify' (addInlineExpr (funName fun) (IEKeep t))
      Nothing -> do
        exceptIf ("remove function", C.posOf r)

    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} $
        ctx'{returnType = rt}

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

  -- Type definitions
  C.CDeclExt (C.CDecl (C.CStorageSpec (C.CTypedef nif) : rst) [item] ni) -> runMaybeT do
    (ix, dd) <- case item of
      C.CDeclarationItem (C.CDeclr (Just ix) dd Nothing _ _) Nothing Nothing ->
        pure (ix, dd)
      i -> notSupportedYet (void i) ni

    modify' (addTypeDef ix ITDelete)

    keep <- reduceStructDeclaration rst
    (bt, rst') <- updateCDeclarationSpecifiers keepAll rst

    (t, _) <- updateCDerivedDeclarators bt (repeat True) dd

    unless keep do
      modify' (addTypeDef ix (ITInline t rst'))
      exceptIf ("inline typedef " <> C.identToString ix, C.posOf ni)

    modify' (addTypeDef ix (ITKeep t))
    pure $ C.CDeclExt $ C.CDecl (C.CStorageSpec (C.CTypedef nif) : rst') [item] ni

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

    markDeleted items

    -- TODO: Actually we should split it up here
    let isStatic = flip any items \case
          (C.CDeclarationItem (C.CDeclr (Just fid) (C.CFunDeclr{} : _) _ _ _) _ _) -> do
            maybe True funIsStatic (lookupFunction ctx fid)
          _ow -> True

    keep <- reduceStructDeclaration spec

    (bt, spec') <- updateCDeclarationSpecifiers keepAll{sfKeepStatic = isStatic} spec
    -- Try to remove each declaration item
    items' <-
      flip collect items \case
        di@(C.CDeclarationItem (C.CDeclr mid dd Nothing [] ni2) einit size) -> do
          case dd of
            C.CFunDeclr{} : _ -> do
              mfun <- case mid of
                Just fid ->
                  Just <$> liftMaybe (lookupFunction ctx fid)
                Nothing ->
                  pure Nothing
              let ff = fromMaybe (repeat True) (mfun >>= funParams)
              (nonVoid -> t, dd') <-
                evalStateT (updateCDerivedDeclarators bt ff dd) ctx
              case mid of
                Just fid -> do
                  modify' (addInlineExpr fid IEDelete)
                  exceptIf ("remove function declaration", C.posOf ni2)
                  modify' (addInlineExpr fid (IEKeep t))
                Nothing -> do
                  exceptIf ("remove function", C.posOf ni2)
              pure (C.CDeclarationItem (C.CDeclr mid dd' Nothing [] ni2) einit size)
            _dd -> reduceCDeclarationItem bt di
        a -> notSupportedYet (a $> ()) ni

    -- Somtimes we just declare a struct or a typedef.
    when (not keep && List.null items') do
      guard (AllowEmptyDeclarations `isIn` ctx || List.null items)
      exceptIf ("remove declaration", C.posOf ni)

    pure $ C.CDeclExt $ C.CDecl spec' items' ni
  _r -> notSupportedYet' r

wrapCCompound :: C.CStatement C.NodeInfo -> C.CStatement C.NodeInfo
wrapCCompound = \case
  s@(C.CCompound{}) -> s
  s -> C.CCompound [] [C.CBlockStmt s] C.undefNode

isStaticFromSpecs :: [C.CDeclarationSpecifier C.NodeInfo] -> Bool
isStaticFromSpecs = any \case
  (C.CStorageSpec (C.CStatic _)) -> True
  _ow -> False

{- | This checks the current declaration and reduces any new struct found here.
Returns true if the specifier is requried.
-}
reduceStructDeclaration
  :: ( MonadReduce Lab m
     , MonadState Context m
     , MonadPlus m
     )
  => [C.CDeclarationSpecifier C.NodeInfo]
  -> m Bool
reduceStructDeclaration =
  fmap or . mapM \case
    C.CTypeSpec (C.CEnumType (C.CEnum mid mf _ _) ni) -> do
      case mid of
        Just eid -> do
          case mf of
            Just times -> forM_ times \(C.CEnumVar ix _) -> do
              modify' (addInlineExpr ix IEDelete)
            Nothing -> pure ()
          modify' (addEnum eid INDelete)
          exceptIf ("delete enum " <> C.identToString eid, C.posOf ni)
          modify' (addEnum eid INKeep)
          case mf of
            Just times -> forM_ times \(C.CEnumVar ix _) -> do
              modify' (addInlineExpr ix (IEKeep TNum))
            Nothing -> pure ()

          pure True
        Nothing -> do
          pure False
    C.CTypeSpec (C.CSUType (C.CStruct tag mid mfields _ _) ni) -> case mid of
      Just sid -> do
        struct <- gets (Map.lookup sid . structs)
        let reduce fields = do
              exceptIf ("remove struct " <> C.identToString sid, C.posOf ni)
              modify' (addStruct sid (ISDeclared tag))
              (ft, _) <- mapAndUnzipM (structField sid) fields
              modify' (addStruct sid (ISKeep (StructType tag (Just sid) (concat ft))))
              pure True
        case struct of
          Just (ISDeclared _) ->
            case mfields of
              Just fields -> reduce fields
              Nothing -> pure False
          Just (ISKeep _) -> do
            pure False
          Just ISDelete -> do
            case mfields of
              Just fields -> reduce fields
              Nothing -> pure True
          Nothing -> do
            modify' (addStruct sid ISDelete)
            case mfields of
              Just fields -> reduce fields
              Nothing -> do
                exceptIf ("remove struct declaration " <> C.identToString sid, C.posOf ni)
                modify' (addStruct sid (ISDeclared tag))
                pure True
      Nothing -> pure False
    _ow -> pure False
 where
  structField sid = \case
    C.CDecl spec items ni -> do
      res <- runMaybeT $ updateCDeclarationSpecifiers keepAll spec
      case res of
        Just (bt, spec') -> do
          (fields, items') <- flip mapAndUnzipM items \case
            (C.CDeclarationItem (C.CDeclr mid dd Nothing [] ni3) ini ni2) -> do
              let fid = fromMaybe (error "all struct fields should be named") mid
              res' <- runMaybeT $ do
                (nonVoid -> t, dd') <- updateCDerivedDeclarators bt (repeat True) dd
                exceptIf ("remove field " <> C.identToString sid <> "." <> C.identToString fid, C.posOf ni)
                pure (t, dd')
              case res' of
                Nothing -> pure ((fid, Nothing), Nothing)
                Just (t, dd') -> pure ((fid, Just t), Just $ C.CDeclarationItem (C.CDeclr mid dd' Nothing [] ni3) ini ni2)
            a -> notSupportedYet a ni
          case catMaybes items' of
            [] -> pure (fields, Nothing)
            items'' -> pure (fields, Just (C.CDecl spec' items'' ni))
        Nothing ->
          pure
            ( map (\i -> (fromMaybe (error "all struct fields should be named") (name i), Nothing)) items
            , Nothing
            )
    a@(C.CStaticAssert{}) -> notSupportedYet' a

reduceCDeclarationItem
  :: ( MonadReduce Lab m
     , MonadState Context m
     , MonadPlus m
     )
  => Voidable
  -> C.CDeclarationItem C.NodeInfo
  -> m (C.CDeclarationItem C.NodeInfo)
reduceCDeclarationItem bt = \case
  di@(C.CDeclarationItem (C.CDeclr mid dd Nothing [] ni) einit Nothing) -> do
    ctx <- get
    case mid of
      Just vid -> do
        (nonVoid -> t, dd') <- evalStateT (updateCDerivedDeclarators bt (repeat True) dd) ctx
        einit' <- case einit of
          Just einit2 -> do
            (einit', inlinable) <- reduceCInitializer t einit2 ctx
            case inlinable of
              Just e' -> do
                modify' (addInlineExpr vid (IEInline e'))
                exceptIf ("inline variable " <> C.identToString vid, C.posOf ni)
              Nothing -> do
                exceptIf ("delete variable", C.posOf ni)
            pure (Just einit')
          Nothing -> do
            exceptIf ("delete uninitilized variable", C.posOf ni)
            whenSplit
              (t == TNum)
              ("initilize variable", C.posOf ni)
              (pure . Just $ C.CInitExpr zeroExpr C.undefNode)
              (pure Nothing)
        modify' (addInlineExpr vid (IEKeep t))
        let decl' = C.CDeclr mid dd' Nothing [] ni
        pure (C.CDeclarationItem decl' einit' Nothing)
      Nothing -> do
        exceptIf ("remove unnamed declaration item", C.posOf ni)
        pure di
  a -> notSupportedYet a C.undefNode

reduceCInitializer
  :: (MonadReduce Lab m)
  => Type
  -> C.CInitializer C.NodeInfo
  -> Context
  -> m (C.CInitializer C.NodeInfo, Maybe C.CExpr)
reduceCInitializer t einit ctx = case einit of
  C.CInitExpr e ni2 -> do
    e' <- fromMaybe (pure zeroExpr) $ reduceCExpr e (exactly t) ctx
    pure
      ( C.CInitExpr e' ni2
      , case e' of
          C.CConst _ -> Just e'
          C.CVar _ _ -> Just e'
          _ow -> Nothing
      )
  C.CInitList (C.CInitializerList items) ni2 -> do
    items' <- case t of
      TStruct stct -> do
        let fields = fieldsOfStruct ctx stct
        let i'' = catMaybes $ zipWith (\(_, t') i -> (i,) <$> t') fields items
        forM i'' \((p, r), t') -> do
          (r', _) <- reduceCInitializer t' r ctx
          pure (p, r')
      TPointer (NonVoid t') -> do
        forM items \(p, r) -> do
          (r', _) <- reduceCInitializer t' r ctx
          pure (p, r')
      _ow ->
        -- "Unexpected type of init list: " <> show t <> " at " <> show (C.posOf ni2)
        pure items
    pure (C.CInitList (C.CInitializerList items') ni2, Nothing)

reduceCCompoundBlockItem
  :: (MonadReduce Lab m, HasCallStack)
  => StmtContext
  -> C.CCompoundBlockItem C.NodeInfo
  -> StateT Context m [C.CCompoundBlockItem C.NodeInfo]
reduceCCompoundBlockItem lab r = do
  case r of
    C.CBlockStmt smt -> do
      ctx <- get
      msmt <- runMaybeT $ reduceCStatement smt lab ctx
      case msmt of
        Just smt' -> do
          case smt' of
            C.CCompound [] ss _ ->
              whenSplit
                (all (\case C.CBlockStmt _ -> True; _ow -> False) ss)
                ("expand compound statment", C.posOf r)
                (pure ss)
                (pure [C.CBlockStmt smt'])
            _ow -> pure [C.CBlockStmt smt']
        Nothing -> pure []
    C.CBlockDecl (C.CDecl spec items ni) -> fmap (fromMaybe []) . runMaybeT $ do
      ctx <- get

      markDeleted items

      keep <- reduceStructDeclaration spec
      (bt, spec') <- updateCDeclarationSpecifiers keepAll spec

      -- Try to remove each declaration item
      items' <- collect (reduceCDeclarationItem bt) items

      -- Somtimes we just declare a struct or a typedef.
      when (not keep && List.null items') do
        guard (AllowEmptyDeclarations `isIn` ctx || List.null items)
        exceptIf ("remove declaration", C.posOf ni)

      pure [C.CBlockDecl (C.CDecl spec' items' ni)]
    a -> notSupportedYet' a

markDeleted :: (MonadState Context m) => [C.CDeclarationItem C.NodeInfo] -> m ()
markDeleted = mapM_ \case
  C.CDeclarationItem (name -> Just ix) _ _ -> do
    modify' (addInlineExpr ix IEDelete)
  _a -> pure ()

reduceCStatementOrEmptyBlock
  :: (MonadReduce Lab m, HasCallStack)
  => C.CStatement C.NodeInfo
  -> StmtContext
  -> Context
  -> m (C.CStatement C.NodeInfo)
reduceCStatementOrEmptyBlock stmt ids ctx = do
  fromMaybe emptyBlock
    <$> runMaybeT
      ( wrapCCompound <$> reduceCStatement stmt ids ctx
      )

reduceCStatementOrEmptyExpr
  :: (MonadReduce Lab m, HasCallStack)
  => C.CStatement C.NodeInfo
  -> StmtContext
  -> Context
  -> m (C.CStatement C.NodeInfo)
reduceCStatementOrEmptyExpr stmt ids ctx = do
  fromMaybe (C.CExpr Nothing C.undefNode)
    <$> runMaybeT (reduceCStatement stmt ids ctx)

emptyBlock :: C.CStatement C.NodeInfo
emptyBlock = C.CCompound [] [] C.undefNode

data StmtContext = StmtContext
  { stmtLabels :: ![C.Ident]
  , stmtInLoop :: !Bool
  }
  deriving (Show, Eq)

etAny :: EType
etAny = EType ETAny False

etNum :: EType
etNum = EType (ETExactly TNum) False

exactly :: Type -> EType
exactly c = EType (ETExactly c) False

-- | Reduce given a list of required labels reduce a c statement, possibly into nothingness.
reduceCStatement
  :: forall m
   . (MonadReduce Lab m, HasCallStack)
  => C.CStatement C.NodeInfo
  -> StmtContext
  -> Context
  -> MaybeT m (C.CStatement C.NodeInfo)
reduceCStatement smt labs ctx = case smt of
  C.CCompound is cbi ni -> do
    cbi' <- lift $ evalStateT (mapM (reduceCCompoundBlockItem labs) cbi) ctx
    pure (C.CCompound is (concat cbi') ni)
  C.CWhile e s dow ni -> split
    ("remove while loop", C.posOf ni)
    do
      reduceCStatement s labs ctx
    do
      s' <- reduceCStatement s labs{stmtInLoop = True} ctx
      e' <- fromMaybe (pure zeroExpr) (reduceCExpr e etNum ctx)
      pure $ C.CWhile e' s' dow ni
  C.CExpr me ni -> do
    case me of
      Just e -> do
        if DoNoops `isIn` ctx
          then do
            e' <-
              maybeSplit ("change to noop", C.posOf smt) $
                reduceCExpr e etAny ctx
            pure $ C.CExpr e' ni
          else do
            re' <- liftMaybe $ reduceCExpr e etAny ctx
            exceptIf ("remove expr statement", C.posOf smt)
            e' <- re'
            pure $ C.CExpr (Just e') ni
      Nothing -> do
        exceptIf ("remove expr statement", C.posOf smt)
        pure $ C.CExpr Nothing ni
  C.CReturn me ni -> do
    re :: m (Maybe C.CExpr) <- case me of
      Just e -> do
        case returnType ctx of
          NonVoid rt -> do
            res :: (m C.CExpr) <- liftMaybe (reduceCExpr e (exactly rt) ctx)
            pure (Just <$> res)
          Void -> pure (pure Nothing)
      Nothing -> pure (pure Nothing)
    exceptIf ("remove return statement", C.posOf smt)
    e <- lift re
    pure $ C.CReturn e ni
  C.CIf e s els ni -> do
    e' <- maybeSplit ("remove condition", C.posOf e) $ reduceCExpr e etNum ctx
    els' <- lift . runMaybeT $ do
      els' <- liftMaybe els
      exceptIf ("remove else branch", C.posOf e)
      reduceCStatement els' labs ctx
    ms' <- lift . runMaybeT $ do
      exceptIf ("remove if branch", C.posOf e)
      reduceCStatement s labs ctx
    case (e', ms', els') of
      (Nothing, Nothing, Nothing) -> empty
      (Just e'', Just s', Nothing) -> pure $ C.CIf e'' s' Nothing ni
      (Nothing, Just s', Just x) -> pure $ C.CIf zeroExpr s' (Just x) ni
      (Just e'', Just s', Just x) -> pure $ C.CIf e'' s' (Just x) ni
      (Just e'', Nothing, Nothing) -> pure $ C.CExpr (Just e'') C.undefNode
      (Nothing, Nothing, Just x) -> pure x
      (Just e'', Nothing, Just x) -> pure $ C.CIf e'' emptyBlock (Just x) ni
      (Nothing, Just s', Nothing) -> pure s'
  C.CFor e1 e2 e3 s ni -> case e1 of
    C.CForDecl d@(C.CDecl spec items ni') -> split
      ("remove the for loop", C.posOf ni)
      (reduceCStatement (C.CCompound [] [C.CBlockDecl d, C.CBlockStmt s] C.undefNode) labs ctx)
      do
        (bt, spec') <- evalStateT (updateCDeclarationSpecifiers keepAll spec) ctx
        (items', ctx') <- flip runStateT ctx do
          markDeleted items
          collect (reduceCDeclarationItem bt) items
        e2' <- runMaybeT do
          e2' <- liftMaybe e2
          re2' <- liftMaybe (reduceCExpr e2' etAny ctx')
          exceptIf ("remove check", C.posOf e2')
          re2'
        e3' <- runMaybeT do
          e3' <- liftMaybe e3
          re3' <- liftMaybe (reduceCExpr e3' etAny ctx')
          exceptIf ("remove iterator", C.posOf e3')
          re3'
        let e2'' =
              if AllowInfiniteForLoops `isIn` ctx || isNothing e2
                then e2'
                else e2' <|> Just zeroExpr
        s' <- reduceCStatementOrEmptyExpr s labs{stmtInLoop = True} ctx'
        pure $ C.CFor (C.CForDecl (C.CDecl spec' items' ni')) e2'' e3' s' ni
    C.CForInitializing e -> split
      ("remove the for loop", C.posOf ni)
      ( reduceCStatement
          ( C.CCompound
              []
              [C.CBlockStmt (C.CExpr e C.undefNode), C.CBlockStmt s]
              C.undefNode
          )
          labs
          ctx
      )
      do
        e' <-
          maybeSplit ("remove initializer", C.posOf ni) $
            e >>= \e' ->
              reduceCExpr e' etAny ctx
        e2' <- runMaybeT do
          e2' <- liftMaybe e2
          re2' <- liftMaybe (reduceCExpr e2' etNum ctx)
          exceptIf ("remove check", C.posOf e2')
          re2'
        e3' <- runMaybeT do
          e3' <- liftMaybe e3
          re3' <- liftMaybe (reduceCExpr e3' etAny ctx)
          exceptIf ("remove iterator", C.posOf e3')
          re3'
        let e2'' =
              if AllowInfiniteForLoops `isIn` ctx || isNothing e2
                then e2'
                else e2' <|> Just zeroExpr
        s' <- reduceCStatementOrEmptyExpr s labs{stmtInLoop = True} ctx
        pure $ C.CFor (C.CForInitializing e') e2'' e3' s' ni
    d -> notSupportedYet d ni
  C.CLabel i s [] ni -> do
    if i `List.elem` stmtLabels labs
      then do
        s' <- lift $ reduceCStatementOrEmptyExpr s labs ctx
        pure $ C.CLabel i s' [] ni
      else do
        empty
  C.CGoto i ni ->
    if i `List.elem` stmtLabels labs
      then do
        exceptIf ("remove goto", C.posOf smt)
        pure $ C.CGoto i ni
      else empty
  C.CBreak n ->
    if stmtInLoop labs
      then do
        exceptIf ("remove break", C.posOf smt)
        pure $ C.CBreak n
      else empty
  C.CCont n ->
    if stmtInLoop labs
      then do
        exceptIf ("remove continue", C.posOf smt)
        pure $ C.CCont n
      else empty
  a -> notSupportedYet' a

-- | If the condition is statisfied try to reduce to the a.
whenSplit :: (MonadReduce Lab m) => Bool -> Lab -> m a -> m a -> m a
whenSplit cn lab a b
  | cn = split lab a b
  | otherwise = b

maybeSplit :: (MonadReduce Lab m) => Lab -> Maybe (m a) -> m (Maybe a)
maybeSplit lab = \case
  Just r -> do
    split lab (pure Nothing) (Just <$> r)
  Nothing -> do
    pure Nothing

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

-- | The expected type
data EType = EType
  { etSet :: !ETSet
  , etAssignable :: !Bool
  }
  deriving (Show, Eq)

data ETSet
  = ETExactly !Type
  | ETStructWithField !C.Ident !ETSet
  | ETPointer !ETSet
  | ETAny
  deriving (Show, Eq)

checkExpectedType :: (MonadPlus m) => Context -> Voidable -> EType -> m ()
checkExpectedType ctx (NonVoid t) et = guard $ isExpectedType ctx t et
checkExpectedType _ Void _ = pure ()

match :: Type -> Type -> Bool
match = curry \case
  (TPointer Void, TPointer _) -> True
  (TPointer _, TPointer Void) -> True
  (TPointer (NonVoid a), TPointer (NonVoid b)) -> a `match` b
  (t1, t2) -> t1 == t2

isExpectedType :: Context -> Type -> EType -> Bool
isExpectedType ctx = \c et ->
  -- pTraceWith (\a -> "check " <> show a <> " " <> show c <> " " <> show et) $
  go c (etSet et)
 where
  go c = \case
    ETExactly t -> t `match` c
    ETAny -> True
    ETStructWithField ix et -> case c of
      TStruct s -> fromMaybe False do
        let fields = fieldsOfStruct ctx s
        (_, mt) <- liftMaybe $ List.find (\(a, _) -> ix == a) fields
        t' <- liftMaybe mt
        pure $ go t' et
      _ow -> False
    ETPointer t' ->
      case c of
        TPointer Void -> True
        TPointer (NonVoid c') -> go c' t'
        _ow -> False

fieldsOfStruct :: (HasCallStack) => Context -> Either C.Ident StructType -> [(C.Ident, Maybe Type)]
fieldsOfStruct ctx (Left ix) =
  case lookupStruct ctx ix of
    ISKeep a -> structTypeFields a
    _ow -> error ("Something bad happend")
fieldsOfStruct _ (Right a) = structTypeFields a

etUnPointer :: EType -> Maybe EType
etUnPointer t =
  -- pTraceWith (\t' -> "unpoint " <> show t <> " " <> show t') $
  case etSet t of
    ETPointer t' -> Just t{etSet = t'}
    ETExactly (TPointer Void) -> Just t{etSet = ETAny}
    ETExactly (TPointer (NonVoid t')) -> Just t{etSet = ETExactly t'}
    _ow -> Nothing

checkNotAssignable :: (MonadPlus m) => EType -> m ()
checkNotAssignable = guard . not . etAssignable

msplit :: (MonadReduce Lab m) => Lab -> Maybe (m a) -> Maybe (m a) -> Maybe (m a)
msplit l m1 m2 = do
  case m1 of
    Just a -> Just $ case m2 of
      Just b -> split l a b
      Nothing -> a
    Nothing -> m2
{-# INLINE msplit #-}

inferType :: Context -> C.CExpr -> Maybe Voidable
inferType ctx = \case
  C.CVar i _ -> do
    case lookupVariable ctx i of
      IEInline e -> inferType ctx e
      IEKeep t -> pure (NonVoid t)
      IEDelete -> Nothing
  C.CUnary i e _ -> do
    t <- inferType ctx e
    case i of
      C.CIndOp -> case t of
        NonVoid (TPointer t') -> pure t'
        Void -> pure Void
        _ow -> Nothing
      C.CAdrOp -> pure (NonVoid (TPointer t))