Optimization.hs

{-# LANGUAGE GADTs, DataKinds, TypeApplications, RankNTypes, ScopedTypeVariables, ConstraintKinds #-}
{-# LANGUAGE DeriveAnyClass, DeriveGeneric #-}
{-# LANGUAGE FlexibleInstances, FlexibleContexts #-}
{-# LANGUAGE QuasiQuotes, RecordWildCards, TupleSections, LambdaCase #-}

module Clang.Coformat.Optimization where

import qualified Data.ByteString.Lazy.Char8 as BSL
import qualified Data.HashMap.Strict as HM
import qualified Data.Text as T
import qualified Data.Text.Lazy.Encoding as TL
import Control.Concurrent.Async.Pool
import Control.Lens
import Control.Monad.Except.CoHas
import Control.Monad.Extra
import Control.Monad.Logger
import Control.Monad.Reader.Has hiding(update)
import Control.Monad.State.Strict
import Data.Foldable
import Data.List
import Data.Ord
import Data.Proxy
import Data.String.Interpolate.IsString
import Numeric.Natural
import System.Command.QQ

import Clang.Coformat.Score
import Clang.Coformat.StyOpts
import Clang.Coformat.Util
import Clang.Coformat.Variables
import Clang.Format.Descr

data FmtEnv = FmtEnv
  { baseStyle :: T.Text
  , preparedFiles :: [PreparedFile]
  , constantOpts :: [ConfigItemT 'Value]
  }

data OptEnv = OptEnv
  { categoricalVariables :: [IxedCategoricalVariable]
  , integralVariables :: [IxedIntegralVariable]
  , maxSubsetSize :: Natural
  }

runClangFormat :: (MonadError err m, CoHas UnexpectedFailure err, CoHas ExpectedFailure err, MonadIO m, MonadLogger m)
               => PreparedFile -> String -> BSL.ByteString -> m Score
runClangFormat prepared logStr formattedSty = do
  stdout <- checked [sh|clang-format --style='#{unpackedSty}' #{path}|]
  let dist = calcScore prepared $ BSL.toStrict $ TL.encodeUtf8 stdout
  logDebugN [i|#{logStr}: #{dist}|]
  pure dist
  where
    path = filename prepared
    unpackedSty = BSL.unpack formattedSty

type OptMonad err r m = (MonadLoggerIO m, MonadError err m, CoHas UnexpectedFailure err, MonadReader r m, Has FmtEnv r)

runClangFormatFiles :: (OptMonad err r m, CoHas ExpectedFailure err)
                    => [ConfigItemT 'Value] -> String -> m Score
runClangFormatFiles varOpts logStr = do
  FmtEnv { .. } <- ask
  let sty = StyOpts { basedOnStyle = baseStyle, additionalOpts = constantOpts <> varOpts }
  let formattedSty = formatStyArg sty
  fmap mconcat $ forM preparedFiles $ \prepared -> runClangFormat prepared [i|#{logStr} at #{filename prepared}|] formattedSty

chooseBaseStyle :: (MonadError String m, MonadLoggerIO m) => [T.Text] -> [ConfigItemT 'Value] -> [PreparedFile] -> m (T.Text, Score)
chooseBaseStyle baseStyles predefinedOpts files = do
  sty2dists <- forConcurrently' ((,) <$> baseStyles <*> files) $ \(sty, file) -> do
    let formattedArg = formatStyArg StyOpts { basedOnStyle = sty, additionalOpts = predefinedOpts }
    convert (show @(Either ExpectedFailure UnexpectedFailure)) $ (sty,) <$> runClangFormat file [i|Initial guess for #{sty} at #{filename file}|] formattedArg
  let accumulated = HM.toList $ HM.fromListWith (<>) sty2dists
  forM_ accumulated $ \(sty, acc) -> logInfoN [i|Initial accumulated guess for #{sty}: #{acc}|]
  pure $ minimumBy (comparing snd) accumulated

variateAt :: forall a. (Variate a, Foldable (VariateResult a))
          => Proxy a -> Int -> [ConfigItemT 'Value] -> [[ConfigItemT 'Value]]
variateAt _ idx opts = [ update idx (updater v') opts | v' <- toList variated ]
  where
    thePrism :: Prism' (ConfigTypeT 'Value) a
    thePrism = varPrism
    variated = variate $ value (opts !! idx) ^?! thePrism
    updater v cfg = cfg { value = value cfg & thePrism .~ v }

data OptState = OptState
  { currentOpts :: [ConfigItemT 'Value]
  , currentScore :: Score
  } deriving (Show)

initOptState :: [ConfigItemT 'Value] -> Score -> OptState
initOptState currentOpts currentScore = OptState { .. }

dropExpectedFailures :: OptMonad err r m
                     => (forall err' r' m'. (OptMonad err' r' m', CoHas ExpectedFailure err') => m' Score)
                     -> m Score
dropExpectedFailures act = do
  res <- runExceptT act
  case res of
       Left (UnexpectedFailure failure) -> throwError failure
       Left (ExpectedFailure failure) -> do
         logErrorN [i|Unable to run the formatter: #{show failure}|]
         pure maxBound
       Right sc -> pure sc

variateSubset :: [SomeIxedVariable] -> [ConfigItemT 'Value] -> [[ConfigItemT 'Value]]
variateSubset [] opts = [opts]
variateSubset (SomeIxedVariable (IxedVariable (MkDV (_ :: a)) idx) : rest) opts = concatMap (variateSubset rest) $ variateAt @a Proxy idx opts

showVariated :: [SomeIxedVariable] -> [ConfigItemT 'Value] -> String
showVariated vars opts = intercalate ", " [showVar var | var <- vars]
  where
    showVar (SomeIxedVariable (IxedVariable _ idx)) = [i|#{name $ opts !! idx} -> #{value $ opts !! idx}|]

chooseBestSubset :: (OptMonad err r m, Has OptState r, Has TaskGroup r)
                 => Natural -> [SomeIxedVariable] -> m (Maybe ([ConfigItemT 'Value], Score))
chooseBestSubset subsetSize ixedVariables = do
  OptState { .. } <- ask
  partialResults <- forConcurrentlyPooled (subsetsN subsetSize ixedVariables) $ \someVarsSubset -> do
    opt2scores <- forM (variateSubset someVarsSubset currentOpts) $ \opts' ->
      fmap (opts',) $ dropExpectedFailures $ runClangFormatFiles opts' $ showVariated someVarsSubset opts'
    let (bestOpts, bestScore) = minimumBy (comparing snd) opt2scores
    when (bestScore < currentScore) $
      logInfoN [i|Total dist for #{showVariated someVarsSubset bestOpts}: #{currentScore} -> #{bestScore}|]
    pure (bestOpts, bestScore, someVarsSubset)
  let (bestOpts, bestScore, bestVarsSubset) = minimumBy (comparing (^. _2)) partialResults
  if bestScore < currentScore
    then do
      logInfoN [i|Choosing #{showVariated bestVarsSubset bestOpts}|]
      pure $ Just (bestOpts, bestScore)
    else pure Nothing

stepGDGeneric' :: (OptMonad err r m, Has TaskGroup r, Has OptEnv r, MonadState OptState m)
               => Natural -> [OptEnv -> [SomeIxedVariable]] -> m ()
stepGDGeneric' subsetSize varGetters = do
  current <- get
  fmtEnv@FmtEnv { .. } <- ask
  optEnv <- ask
  tg :: TaskGroup <- ask
  runReaderT (chooseBestSubset subsetSize $ concatMap ($ optEnv) varGetters) (current, fmtEnv, tg) >>=
    \case Nothing -> pure ()
          Just (opts', score') -> do
            logInfoN [i|Total score after optimization on all files: #{score'}|]
            put OptState { currentOpts = opts', currentScore = score' }

stepGDGeneric :: (OptMonad err r m, Has TaskGroup r, Has OptEnv r, MonadState OptState m)
              => Natural -> [OptEnv -> [SomeIxedVariable]] -> m ()
stepGDGeneric subsetSize varGetters = whenM ((> mempty) <$> gets currentScore) $ stepGDGeneric' subsetSize varGetters

fixGD :: (OptMonad err r m, Has TaskGroup r, Has OptEnv r, MonadState OptState m, err ~ UnexpectedFailure)
      => Maybe Int -> Natural -> m ()
fixGD (Just 0) _ = pure ()
fixGD counter curSubsetSize = do
  maxSubsetSize' <- asks maxSubsetSize
  if curSubsetSize > maxSubsetSize'
    then logInfoN [i|Done optimizing|]
    else do
      startScore <- gets currentScore
      stepGDGeneric curSubsetSize [asSome . categoricalVariables, asSome . integralVariables]
      endScore <- gets currentScore
      logInfoN [i|Full optimization step done, went from #{startScore} to #{endScore}|]
      if startScore /= endScore
        then fixGD (subtract 1 <$> counter) 1
        else do
          logInfoN [i|Done optimizing with subset size #{curSubsetSize}, stopped at score #{endScore}|]
          fixGD (subtract 1 <$> counter) (curSubsetSize + 1)