ReduceC.hs

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE NoMonomorphismRestriction #-}

module ReduceC (
  defaultReduceC,
  -- reduceCTranslUnit,

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

  -- * Helpers
  prettyIdent,
) where

import Control.Monad.Reduce
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 qualified Data.Set as Set
import Data.Vector.Internal.Check (HasCallStack)

-- import Debug.Trace

import qualified Control.Monad.IRTree as IRTree
import Data.Monoid
import qualified Language.C as C
import qualified Language.C.Data.Ident as C

data Context = Context
  { keywords :: !(Set.Set Keyword)
  , typeDefs :: !(Map.Map C.Ident InlineType)
  , inlineExprs :: !(Map.Map C.Ident InlineExpr)
  }
  deriving (Show)

data InlineType
  = ITKeep
  | ITInline ![C.CDeclarationSpecifier C.NodeInfo]
  deriving (Show, Eq)

data InlineExpr
  = IEDelete
  | IEInline !C.CExpr
  | IEKeep
  deriving (Show, Eq)

data Keyword
  = KeepMain
  | DoNoops
  | NoSemantics
  | AllowEmptyDeclarations
  | DisallowVariableInlining
  deriving (Show, Read, Enum, Eq, Ord)

type Lab = (String, C.Position)

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 #-}

addTypeDefs :: [C.Ident] -> InlineType -> Context -> Context
addTypeDefs ids cs Context{..} =
  Context
    { typeDefs =
        foldl' (\a i -> Map.insert i cs a) typeDefs ids
    , ..
    }

addInlineExpr :: C.Ident -> InlineExpr -> Context -> Context
addInlineExpr i e Context{..} =
  Context
    { inlineExprs = Map.insert i e inlineExprs
    , ..
    }

addKeyword :: Keyword -> Context -> Context
addKeyword k Context{..} =
  Context
    { keywords = Set.insert k keywords
    , ..
    }

-- deleteKeyword :: Keyword -> Context -> Context
-- deleteKeyword k Context{..} =
--   Context
--     { keywords = Set.delete k keywords
--     , ..
--     }

defaultContext :: Context
defaultContext =
  Context
    { keywords = Set.fromList [KeepMain]
    , typeDefs = Map.empty
    , inlineExprs = Map.empty
    }

isIn :: Keyword -> Context -> Bool
isIn k = Set.member k . keywords

prettyIdent :: C.Identifier C.NodeInfo -> [Char]
prettyIdent (C.Ident s _ a) = s ++ " at " ++ show (C.posOfNode a)

reduceCTranslUnit
  :: (MonadReduce Lab m)
  => C.CTranslationUnit C.NodeInfo
  -> Context
  -> m (C.CTranslationUnit C.NodeInfo)
reduceCTranslUnit (C.CTranslUnit es ni) ctx = do
  es' <- foldr reduceCExternalDeclaration (\_ -> pure []) es ctx
  pure $ C.CTranslUnit es' ni

reduceCExternalDeclaration
  :: (MonadReduce Lab m)
  => C.CExternalDeclaration C.NodeInfo
  -> (Context -> m [C.CExternalDeclaration C.NodeInfo])
  -> Context
  -> m [C.CExternalDeclaration C.NodeInfo]
reduceCExternalDeclaration r cont ctx = do
  -- TODO This is slow
  case r of
    C.CFDefExt fun
      | KeepMain `isIn` ctx && maybe False (("main" ==) . C.identToString) (functionName fun) -> do
          r' <- C.CFDefExt <$> reduceCFunDef fun ctx
          (r' :) <$> cont ctx
      | otherwise ->
          case functionName fun of
            Just fid -> do
              split
                ("remove function " <> C.identToString fid, C.posOf r)
                (cont (addInlineExpr fid IEDelete ctx))
                do
                  r' <- C.CFDefExt <$> reduceCFunDef fun ctx
                  (r' :) <$> cont (addInlineExpr fid IEKeep ctx)
            Nothing -> do
              split
                ("remove function", C.posOf r)
                (cont ctx)
                do
                  r' <- C.CFDefExt <$> reduceCFunDef fun ctx
                  (r' :) <$> cont ctx
    C.CDeclExt result ->
      case result of
        -- A typedef
        C.CDecl (C.CStorageSpec (C.CTypedef _) : rst) decl _ -> do
          let [ids] = identifiers decl
          split
            ("inline typedef " <> C.identToString ids, C.posOf r)
            (cont (addTypeDefs [ids] (ITInline rst) ctx))
            ( (C.CDeclExt (inlineTypeDefsCDeclaration result ctx) :)
                <$> cont (addTypeDefs [ids] ITKeep ctx)
            )
        -- A const
        C.CDecl rec decl ni' -> do
          (decl', ctx') <- foldr reduceCDeclarationItem (pure ([], ctx)) decl
          case decl' of
            []
              | AllowEmptyDeclarations `isIn` ctx' ->
                  split ("remove empty declaration", C.posOf r) (cont ctx') do
                    (C.CDeclExt (inlineTypeDefsCDeclaration (C.CDecl rec decl' ni') ctx) :) <$> cont ctx'
              | otherwise -> cont ctx'
            _ow -> (C.CDeclExt (inlineTypeDefsCDeclaration (C.CDecl rec decl' ni') ctx) :) <$> cont ctx'
        a -> don'tHandle a
    _r -> don'tHandle r

reduceCFunDef
  :: (MonadReduce Lab m, HasCallStack)
  => C.CFunctionDef C.NodeInfo
  -> Context
  -> m (C.CFunctionDef C.NodeInfo)
reduceCFunDef (C.CFunDef spc dec cdecls smt ni) ctx = do
  smt' <- reduceCStatementOrEmptyBlock smt ctx'
  pure $
    C.CFunDef
      (inlineTypeDefsSpecs spc ctx)
      (inlineTypeDefsCDeclarator dec ctx)
      (map (`inlineTypeDefsCDeclaration` ctx) cdecls)
      smt'
      ni
 where
  !ctx' = foldr (`addInlineExpr` IEKeep) ctx ids
  ids = params dec

params :: C.CDeclarator C.NodeInfo -> [C.Ident]
params = \case
  C.CDeclr _ [C.CFunDeclr (C.CFunParamsNew decls _) _ _] _ _ _ ->
    decls & concatMap \case
      C.CDecl _ items _ ->
        items & concatMap \case
          C.CDeclarationItem (C.CDeclr (Just idx) _ _ _ _) _ _ -> [idx]
          _ow -> []
      a -> don'tHandle a
  a -> don'tHandle a

reduceCCompoundBlockItem
  :: (MonadReduce Lab m, HasCallStack)
  => C.CCompoundBlockItem C.NodeInfo
  -> (Context -> m [C.CCompoundBlockItem C.NodeInfo])
  -> Context
  -> m [C.CCompoundBlockItem C.NodeInfo]
reduceCCompoundBlockItem r cont ctx = do
  case r of
    C.CBlockStmt smt -> do
      case reduceCStatement smt ctx of
        Just rsmt -> split ("remove statement", C.posOf r) (cont ctx) do
          smt' <- rsmt
          case smt' of
            C.CCompound [] ss _ -> do
              split ("expand compound statment", C.posOf r) ((ss <>) <$> cont ctx) do
                (C.CBlockStmt smt' :) <$> cont ctx
            _ow -> do
              (C.CBlockStmt smt' :) <$> cont ctx
        Nothing -> cont ctx
    C.CBlockDecl declr -> do
      case declr of
        C.CDecl rec decl ni' -> do
          (decl', ctx') <- foldr reduceCDeclarationItem (pure ([], ctx)) decl
          case decl' of
            []
              | AllowEmptyDeclarations `isIn` ctx' ->
                  split ("remove empty declaration", C.posOf r) (cont ctx') do
                    (C.CBlockDecl (inlineTypeDefsCDeclaration (C.CDecl rec decl' ni') ctx) :) <$> cont ctx'
              | otherwise -> cont ctx'
            _ow -> (C.CBlockDecl (inlineTypeDefsCDeclaration (C.CDecl rec decl' ni') ctx) :) <$> cont ctx'
        d -> don'tHandle d
    a -> don'tHandle a

reduceCDeclarationItem
  :: (MonadReduce Lab m)
  => C.CDeclarationItem C.NodeInfo
  -> m ([C.CDeclarationItem C.NodeInfo], Context)
  -> m ([C.CDeclarationItem C.NodeInfo], Context)
reduceCDeclarationItem d ma = case d of
  C.CDeclarationItem
    dr@(C.CDeclr (Just i) [] Nothing [] ni)
    (Just (C.CInitExpr c ni'))
    Nothing -> do
      (ds, ctx) <- ma
      c' <- fromMaybe (pure zeroExpr) (reduceCExpr c ctx)
      split
        ("inline variable " <> C.identToString i, C.posOf ni)
        (pure (ds, addInlineExpr i (IEInline c') ctx))
        ( pure
            ( inlineTypeDefsCDI (C.CDeclarationItem dr (Just (C.CInitExpr c' ni')) Nothing) ctx : ds
            , addInlineExpr i IEKeep ctx
            )
        )
  C.CDeclarationItem (C.CDeclr (Just i) _ Nothing _ ni) _ Nothing -> do
    (ds, ctx) <- ma
    split
      ("remove variable " <> C.identToString i, C.posOf ni)
      (pure (ds, addInlineExpr i IEDelete ctx))
      (pure (inlineTypeDefsCDI d ctx : ds, addInlineExpr i IEKeep ctx))
  a@(C.CDeclarationItem (C.CDeclr _ _ _ _ ni) _ _) -> do
    don'tHandleWithNodeInfo a ni
  a -> don'tHandle a

reduceCStatementOrEmptyBlock
  :: (MonadReduce Lab m, HasCallStack)
  => C.CStatement C.NodeInfo
  -> Context
  -> m (C.CStatement C.NodeInfo)
reduceCStatementOrEmptyBlock stmt ctx = do
  case reduceCStatement stmt ctx of
    Just ex -> do
      ex
    Nothing -> do
      pure emptyBlock
 where
  emptyBlock = C.CCompound [] [] C.undefNode

reduceCStatement
  :: (MonadReduce Lab m, HasCallStack)
  => C.CStatement C.NodeInfo
  -> Context
  -> Maybe (m (C.CStatement C.NodeInfo))
reduceCStatement smt ctx = case smt of
  C.CCompound is cbi ni -> Just do
    cbi' <- foldr reduceCCompoundBlockItem (\_ -> pure []) cbi ctx
    pure $ C.CCompound is cbi' ni
  C.CWhile e s dow ni -> do
    rs <- reduceCStatement s ctx
    Just do
      e' <- reduceCExprOrZero e ctx
      s' <- rs
      pure $ C.CWhile e' s' dow ni
  C.CExpr me ni -> do
    case me of
      Just e -> do
        if DoNoops `isIn` ctx
          then Just do
            e' <- maybeSplit ("change to noop", C.posOf smt) $ reduceCExpr e ctx
            pure $ C.CExpr e' ni
          else do
            re <- reduceCExpr e ctx
            Just do
              e' <- re
              pure $ C.CExpr (Just e') ni
      Nothing ->
        Just $ pure $ C.CExpr Nothing ni
  C.CReturn me ni -> Just do
    case me of
      Just e -> do
        e' <- reduceCExprOrZero e ctx
        pure $ C.CReturn (Just e') ni
      Nothing ->
        pure $ C.CReturn Nothing ni
  C.CIf e s els ni -> Just do
    e' <- maybeSplit ("remove condition", C.posOf e) $ reduceCExpr e ctx
    els' <- case els of
      Just els' -> do
        maybeSplit ("remove else branch", C.posOf els') do
          reduceCStatement els' ctx
      Nothing -> pure Nothing
    s' <- reduceCStatementOrEmptyBlock s ctx
    case (e', els') of
      (Nothing, Nothing) -> pure s'
      (Just e'', Nothing) -> pure $ C.CIf e'' s' Nothing ni
      (Nothing, Just x) -> pure $ C.CIf zeroExpr s' (Just x) ni
      (Just e'', Just x) -> pure $ C.CIf e'' s' (Just x) ni
  C.CFor e1 e2 e3 s ni -> Just $ do
    (me1', ctx') <- case e1 of
      C.CForDecl (C.CDecl rec decl ni') -> do
        (decl', ctx') <- foldr reduceCDeclarationItem (pure ([], ctx)) decl
        res <-
          if null decl'
            then
              whenSplit
                (AllowEmptyDeclarations `isIn` ctx')
                ("remove empty declaration", C.posOf ni')
                (pure Nothing)
                (pure $ Just $ C.CForDecl (C.CDecl rec decl' ni'))
            else pure $ Just $ C.CForDecl (C.CDecl rec decl' ni')
        pure (res, ctx')
      C.CForInitializing e -> do
        e' <- maybeSplit ("remove initializer", C.posOf ni) (e >>= \e' -> reduceCExpr e' ctx)
        whenSplit
          (AllowEmptyDeclarations `isIn` ctx)
          ("remove empty declaration", C.posOf ni)
          (pure (Nothing, ctx))
          (pure (Just $ C.CForInitializing e', ctx))
      d -> don'tHandle d

    s' <- reduceCStatementOrEmptyBlock s ctx'
    case me1' of
      Nothing -> do
        split ("remove the for loop", C.posOf smt) (pure s') do
          e2' <- case e2 of
            Just e2' -> maybeSplit ("remove check", C.posOf e2') (reduceCExpr e2' ctx')
            Nothing -> pure Nothing
          e3' <- case e3 of
            Just e3' -> maybeSplit ("remove iterator", C.posOf e3') (reduceCExpr e3' ctx')
            Nothing -> pure Nothing
          pure $ C.CFor (C.CForInitializing Nothing) e2' e3' s' ni
      Just e1' -> do
        e2' <- case e2 of
          Just e2' -> maybeSplit ("remove check", C.posOf e2') (reduceCExpr e2' ctx')
          Nothing -> pure Nothing
        e3' <- case e3 of
          Just e3' -> maybeSplit ("remove iterator", C.posOf e3') (reduceCExpr e3' ctx')
          Nothing -> pure Nothing
        pure $ C.CFor e1' e2' e3' s' ni
  C.CBreak ni -> Just do
    pure (C.CBreak ni)
  C.CCont ni -> Just do
    pure (C.CCont ni)
  C.CLabel i s [] ni -> Just do
    s' <- reduceCStatementOrEmptyBlock s ctx
    pure $ C.CLabel i s' [] ni
  C.CGoto i ni -> Just do
    pure $ C.CGoto i ni
  a -> don'tHandleWithPos a

--     C.CCompound is cbi ni -> do
--       cbi' <- collect (reduce @C.CCompoundBlockItem) cbi
--       pure $ C.CCompound is cbi' ni
--     C.CExpr e ni -> do
--       e' <- optional do
--         e' <- liftMaybe e
--         reduce @C.CExpression e'
--       pure $ C.CExpr e' ni
--     C.CReturn e ni -> do
--       e' <- traverse (fmap orZero reduce) e
--       pure $ C.CReturn e' ni
--     C.CLabel i s [] ni -> do
--       -- todo fix attrs
--       s' <- reduce s
--       withFallback s' do
--         givenThat (Val.is i)
--         pure $ C.CLabel i s' [] ni
--     C.CWhile e s dow ni -> do
--       e' <- orZero (reduce @C.CExpression e)
--       s' <- reduce s
--       pure $ C.CWhile e' s' dow ni

-- | 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)

reduceCExprOrZero :: (MonadReduce Lab m, HasCallStack) => C.CExpr -> Context -> m C.CExpr
reduceCExprOrZero expr ctx = do
  case reduceCExpr expr ctx of
    Just ex -> do
      split ("replace by zero", C.posOf expr) (pure zeroExpr) ex
    Nothing -> do
      pure zeroExpr
{-# INLINE reduceCExprOrZero #-}

reduceCExpr :: (MonadReduce Lab m, HasCallStack) => C.CExpr -> Context -> Maybe (m C.CExpr)
reduceCExpr expr ctx = case expr of
  C.CBinary o elhs erhs ni -> do
    case reduceCExpr elhs ctx of
      Just elhs' -> case reduceCExpr erhs ctx of
        Just erhs' -> pure do
          split ("reduce to left", C.posOf elhs) elhs' do
            split ("reduce to right", C.posOf erhs) erhs' do
              l' <- elhs'
              r' <- erhs'
              pure $ C.CBinary o l' r' ni
        Nothing ->
          fail "could not reduce right hand side"
      Nothing
        | otherwise -> fail "could not reduce left hand side"
  C.CAssign o elhs erhs ni ->
    case reduceCExpr elhs (addKeyword DisallowVariableInlining ctx) of
      Just elhs' -> case reduceCExpr erhs ctx of
        Just erhs' -> pure do
          split ("reduce to left", C.posOf elhs) elhs' do
            split ("reduce to right", C.posOf erhs) erhs' do
              l' <- elhs'
              r' <- erhs'
              pure $ C.CAssign o l' r' ni
        Nothing ->
          fail "could not reduce right hand side"
      Nothing
        | otherwise -> fail "could not reduce left hand side"
  C.CVar i _ ->
    case Map.lookup i . inlineExprs $ ctx of
      Just mx -> case mx of
        IEKeep -> Just (pure expr)
        IEInline mx'
          | DisallowVariableInlining `isIn` ctx -> Nothing
          | otherwise -> Just (pure mx')
        IEDelete ->
          Nothing
      Nothing -> error ("Could not find " <> show i <> " at " <> show (C.posOf expr) <> "\n" <> show (inlineExprs ctx))
  C.CConst x -> Just do
    pure $ C.CConst x
  C.CUnary o elhs ni -> do
    elhs' <- reduceCExpr elhs (addKeyword DisallowVariableInlining ctx)
    Just $ split ("reduce to operant", C.posOf expr) elhs' do
      e <- elhs'
      pure $ C.CUnary o e ni
  C.CCall e es ni -> do
    re <- reduceCExpr e (addKeyword DisallowVariableInlining ctx)
    Just $ do
      e' <- re
      es' <- traverse (`reduceCExprOrZero` ctx) es
      pure $ C.CCall e' es' ni
  C.CCond ec et ef ni -> do
    -- TODO: More fine grained reduction is possible here.
    Just $ do
      ec' <- reduceCExprOrZero ec ctx
      ef' <- reduceCExprOrZero ef ctx
      et' <- case et of
        Just et' -> Just <$> reduceCExprOrZero et' ctx
        Nothing -> pure Nothing
      pure $ C.CCond ec' et' ef' ni
  C.CCast decl e ni -> do
    re <- reduceCExpr e ctx
    Just do
      split ("don't cast", C.posOf ni) re do
        e' <- re
        pure (C.CCast (inlineTypeDefsCDeclaration decl ctx) e' ni)
  C.CIndex e1 e2 ni -> do
    -- TODO: Better reduction is posisble here.
    re1 <- reduceCExpr e1 ctx
    Just do
      e1' <- re1
      e2' <- reduceCExprOrZero e2 ctx
      pure $ C.CIndex e1' e2' ni
  C.CComma items ni -> Just do
    let Just (x, rst) = List.uncons (reverse items)
    rst' <-
      foldr
        ( \e cc -> do
            maybeSplit ("remove expression", C.posOf e) (reduceCExpr e ctx) >>= \case
              Just e' -> (e' :) <$> cc
              Nothing -> cc
        )
        (pure [])
        rst
    x' <- reduceCExprOrZero x ctx
    if List.null rst'
      then pure x'
      else pure $ C.CComma (reverse (x' : rst')) ni
  a -> don'tHandleWithPos a

inlineTypeDefsCDeclaration :: C.CDeclaration C.NodeInfo -> Context -> C.CDeclaration C.NodeInfo
inlineTypeDefsCDeclaration decl ctx =
  case decl of
    C.CDecl items decli ni ->
      C.CDecl (inlineTypeDefsSpecs items ctx) (map (`inlineTypeDefsCDI` ctx) decli) ni
    a -> don'tHandle a

inlineTypeDefsSpecs :: [C.CDeclarationSpecifier C.NodeInfo] -> Context -> [C.CDeclarationSpecifier C.NodeInfo]
inlineTypeDefsSpecs r ctx =
  r & concatMap \case
    a@(C.CTypeSpec (C.CTypeDef idx _)) -> do
      case Map.lookup idx . typeDefs $ ctx of
        Just ITKeep -> [a]
        Just (ITInline res) -> res
        Nothing -> error ("could not find typedef:" <> show idx)
    a -> [a]
{-# NOINLINE inlineTypeDefsSpecs #-}

inlineTypeDefsCDeclarator
  :: C.CDeclarator C.NodeInfo
  -> Context
  -> C.CDeclarator C.NodeInfo
inlineTypeDefsCDeclarator (C.CDeclr idn derivedd st atr ni) ctx =
  C.CDeclr idn (map (inlineTypeDefsX ctx) derivedd) st atr ni

inlineTypeDefsX :: Context -> C.CDerivedDeclarator C.NodeInfo -> C.CDerivedDeclarator C.NodeInfo
inlineTypeDefsX ctx = \case
  C.CFunDeclr (C.CFunParamsNew x y) b c ->
    C.CFunDeclr (C.CFunParamsNew (map (`inlineTypeDefsCDeclaration` ctx) x) y) b c
  C.CArrDeclr a b c -> C.CArrDeclr a b c
  C.CPtrDeclr a b -> C.CPtrDeclr a b
  a -> don'tHandle a

inlineTypeDefsCDI :: C.CDeclarationItem C.NodeInfo -> Context -> C.CDeclarationItem C.NodeInfo
inlineTypeDefsCDI di ctx = case di of
  C.CDeclarationItem a b ni -> C.CDeclarationItem (inlineTypeDefsCDeclarator a ctx) b ni
  a -> don'tHandle a

identifiers :: forall a. (Data a) => a -> [C.Ident]
identifiers d = appEndo (go d) []
 where
  go :: forall a'. (Data a') => a' -> Endo [C.Ident]
  go d' = case cast d' of
    Just l -> Endo (l :)
    Nothing -> gmapQl (<>) mempty go d'

functionName :: C.CFunctionDef C.NodeInfo -> Maybe C.Ident
functionName = \case
  C.CFunDef _ (C.CDeclr ix _ _ _ _) _ _ _ -> ix

don'tHandle :: (HasCallStack, Functor f, Show (f ())) => f C.NodeInfo -> b
don'tHandle f = error (show (f $> ()))

don'tHandleWithPos :: (HasCallStack, Functor f, Show (f ()), C.Pos (f C.NodeInfo)) => f C.NodeInfo -> b
don'tHandleWithPos f = error (show (f $> ()) <> " at " <> show (C.posOf f))

don'tHandleWithNodeInfo :: (HasCallStack, Functor f, Show (f ())) => f C.NodeInfo -> C.NodeInfo -> b
don'tHandleWithNodeInfo f ni = error (show (f $> ()) <> " at " <> show (C.posOf ni))