content: -- Copyright 2023 Google LLC
  --
  -- Use of this source code is governed by a BSD-style
  -- license that can be found in the LICENSE file or at
  -- https://developers.google.com/open-source/licenses/bsd
  
  module PeepholeOptimize (PeepholeOpt (..), peepholeExpr) where
  
  import Data.Word
  import Data.Bits
  import Data.List
  import Data.Bits.Floating
  import GHC.Float
  
  import Types.Core
  import Types.Primitives
  import Name
  import IRVariants
  import qualified Types.OpNames as P
  
  peepholeExpr :: Expr r n -> Expr r n
  peepholeExpr e = case peephole e of
    Just x  -> Atom x
    Nothing -> e
  {-# INLINE peepholeExpr #-}
  
  -- === Peephole optimization = undefined
  
  -- These are context-free (no env!) optimizations of expressions and ops that
  -- are worth doing unconditionally. Builder calls this automatically in `emit`.
  
  class ToExpr e r => PeepholeOpt (e::E) (r::IR) | e -> r where
    peephole :: e n -> Maybe (Atom r n)
  
  instance PeepholeOpt (Expr r) r where
    peephole = \case
      Atom x -> Just x
      PrimOp op -> peephole op
      Project _ i x  -> case x of
        Con con -> Just case con of
          ProdCon xs             -> xs !! i
          DepPair l _ _ | i == 0 -> l
          DepPair _ r _ | i == 1 -> r
          _ -> error "not a product"
        Stuck _ _ -> Nothing
      Unwrap _ x -> case x of
        Con con -> Just case con of
          NewtypeCon _ x' -> x'
          _ -> error "not a newtype"
        Stuck _ _  -> Nothing
      App    _ _ _ -> Nothing
      TabApp _ _ _ -> Nothing
      Case _ _ _   -> Nothing
      TopApp _ _ _ -> Nothing
      Block _ _    -> Nothing
      TabCon _ _ _ -> Nothing
      ApplyMethod _ _ _ _ -> Nothing
    {-# INLINE peephole #-}
  
  instance PeepholeOpt (PrimOp r) r where
    peephole = \case
      MiscOp op -> peephole op
      BinOp op l r -> peepholeBinOp op l r
      _ -> Nothing
    {-# INLINE peephole #-}
  
  peepholeBinOp :: P.BinOp -> Atom r n -> Atom r n -> Maybe (Atom r n)
  peepholeBinOp op x y = case op of
    IAdd -> case (x, y) of
      (Con (Lit x'), y') | getIntLit x' == 0 -> Just y'
      (x', Con (Lit y')) | getIntLit y' == 0 -> Just x'
      (Con (Lit x'), Con (Lit y'))           -> Just $ Con $ Lit $ applyIntBinOp (+) x' y'
      _ -> Nothing
    ISub -> case (x, y) of
      (x', Con (Lit y')) | getIntLit y' == 0 -> Just x'
      (Con (Lit x'), Con (Lit y'))           -> Just $ Con $ Lit $ applyIntBinOp (-) x' y'
      _ -> Nothing
    IMul -> case (x, y) of
      (Con (Lit x'), y') | getIntLit x' == 1 -> Just y'
      (x', Con (Lit y')) | getIntLit y' == 1 -> Just x'
      (Con (Lit x'), Con (Lit y'))           -> Just $ Con $ Lit $ applyIntBinOp (*) x' y'
      _ -> Nothing
    IDiv -> case (x, y) of
      (x', Con (Lit y')) | getIntLit y' == 1 -> Just x'
      _ -> Nothing
    ICmp cop -> case (x, y) of
      (Con (Lit x'), Con (Lit y')) -> Just $ Con $ Lit $ applyIntCmpOp (cmp cop) x' y'
      _ -> Nothing
    FAdd -> case (x, y) of
      (Con (Lit x'), y') | getFloatLit x' == 0 -> Just y'
      (x', Con (Lit y')) | getFloatLit y' == 0 -> Just x'
      (Con (Lit x'), Con (Lit y'))           -> Just $ Con $ Lit $ applyFloatBinOp (+) x' y'
      _ -> Nothing
    FSub -> case (x, y) of
      (x', Con (Lit y')) | getFloatLit y' == 0 -> Just x'
      (Con (Lit x'), Con (Lit y'))           -> Just $ Con $ Lit $ applyFloatBinOp (-) x' y'
      _ -> Nothing
    FMul -> case (x, y) of
      (Con (Lit x'), y') | getFloatLit x' == 1 -> Just y'
      (x', Con (Lit y')) | getFloatLit y' == 1 -> Just x'
      (Con (Lit x'), Con (Lit y'))             -> Just $ Con $ Lit $ applyFloatBinOp (*) x' y'
      _ -> Nothing
    FDiv -> case (x, y) of
      (x', Con (Lit y')) | getFloatLit y' == 1 -> Just x'
      _ -> Nothing
    FCmp cop -> case (x, y) of
      (Con (Lit x'), Con (Lit y')) -> Just $ Con $ Lit $ applyFloatCmpOp (cmp cop) x' y'
      _ -> Nothing
    BOr -> case (x, y) of
      (Con (Lit (Word8Lit x')), Con (Lit (Word8Lit y'))) -> Just $ Con $ Lit $ Word8Lit $ x' .|. y'
      _ -> Nothing
    BAnd -> case (x, y) of
      (Con (Lit (Word8Lit lv)), Con (Lit (Word8Lit rv))) -> Just $ Con $ Lit $ Word8Lit $ lv .&. rv
      _ -> Nothing
    BXor -> Nothing -- TODO
    BShL -> Nothing -- TODO
    BShR -> Nothing -- TODO
    IRem -> Nothing -- TODO
    FPow -> Nothing -- TODO
  {-# INLINE peepholeBinOp #-}
  
  instance PeepholeOpt (MiscOp r) r where
    peephole = \case
      CastOp (TyCon (BaseType (Scalar sTy))) (Con (Lit l)) -> case foldCast sTy l of
        Just l' -> Just $ Con $ Lit l'
        Nothing -> Nothing
      ToEnum ty (Con (Lit (Word8Lit tag))) -> case ty of
        TyCon (SumType cases) -> Just $ Con $ SumCon cases (fromIntegral tag) UnitVal
        _ -> error "Ill typed ToEnum"
      SumTag (Con (SumCon _ tag _)) -> Just $ Con $ Lit $ Word8Lit $ fromIntegral tag
      Select p x y -> case p of
        Con (Lit (Word8Lit p')) -> Just if p' /= 0 then x else y
        _ -> Nothing
      _ -> Nothing
  
  foldCast :: ScalarBaseType -> LitVal -> Maybe LitVal
  foldCast sTy l = case sTy of
    -- TODO: Check that the casts relating to floating-point agree with the
    -- runtime behavior.  The runtime is given by the `ICastOp` case in
    -- ImpToLLVM.hs.  We should make sure that the Haskell functions here
    -- produce bitwise identical results to those instructions, by adjusting
    -- either this or that as called for.
    -- TODO: Also implement casts that may have unrepresentable results, i.e.,
    -- casting floating-point numbers to smaller floating-point numbers or to
    -- fixed-point.  Both of these necessarily have a much smaller dynamic range.
    Int32Type -> case l of
      Int32Lit  _  -> Just l
      Int64Lit  i  -> Just $ Int32Lit  $ fromIntegral i
      Word8Lit  i  -> Just $ Int32Lit  $ fromIntegral i
      Word32Lit i  -> Just $ Int32Lit  $ fromIntegral i
      Word64Lit i  -> Just $ Int32Lit  $ fromIntegral i
      Float32Lit _ -> Nothing
      Float64Lit _ -> Nothing
      PtrLit   _ _ -> Nothing
    Int64Type -> case l of
      Int32Lit  i  -> Just $ Int64Lit  $ fromIntegral i
      Int64Lit  _  -> Just l
      Word8Lit  i  -> Just $ Int64Lit  $ fromIntegral i
      Word32Lit i  -> Just $ Int64Lit  $ fromIntegral i
      Word64Lit i  -> Just $ Int64Lit  $ fromIntegral i
      Float32Lit _ -> Nothing
      Float64Lit _ -> Nothing
      PtrLit   _ _ -> Nothing
    Word8Type -> case l of
      Int32Lit  i  -> Just $ Word8Lit  $ fromIntegral i
      Int64Lit  i  -> Just $ Word8Lit  $ fromIntegral i
      Word8Lit  _  -> Just l
      Word32Lit i  -> Just $ Word8Lit  $ fromIntegral i
      Word64Lit i  -> Just $ Word8Lit  $ fromIntegral i
      Float32Lit _ -> Nothing
      Float64Lit _ -> Nothing
      PtrLit   _ _ -> Nothing
    Word32Type -> case l of
      Int32Lit  i  -> Just $ Word32Lit $ fromIntegral i
      Int64Lit  i  -> Just $ Word32Lit $ fromIntegral i
      Word8Lit  i  -> Just $ Word32Lit $ fromIntegral i
      Word32Lit _  -> Just l
      Word64Lit i  -> Just $ Word32Lit $ fromIntegral i
      Float32Lit _ -> Nothing
      Float64Lit _ -> Nothing
      PtrLit   _ _ -> Nothing
    Word64Type -> case l of
      Int32Lit  i  -> Just $ Word64Lit $ fromIntegral (fromIntegral i :: Word32)
      Int64Lit  i  -> Just $ Word64Lit $ fromIntegral i
      Word8Lit  i  -> Just $ Word64Lit $ fromIntegral i
      Word32Lit i  -> Just $ Word64Lit $ fromIntegral i
      Word64Lit _  -> Just l
      Float32Lit _ -> Nothing
      Float64Lit _ -> Nothing
      PtrLit   _ _ -> Nothing
    Float32Type -> case l of
      Int32Lit  i  -> Just $ Float32Lit $ fixUlp i $ fromIntegral i
      Int64Lit  i  -> Just $ Float32Lit $ fixUlp i $ fromIntegral i
      Word8Lit  i  -> Just $ Float32Lit $ fromIntegral i
      Word32Lit i  -> Just $ Float32Lit $ fixUlp i $ fromIntegral i
      Word64Lit i  -> Just $ Float32Lit $ fixUlp i $ fromIntegral i
      Float32Lit _ -> Just l
      Float64Lit _ -> Nothing
      PtrLit   _ _ -> Nothing
    Float64Type -> case l of
      Int32Lit  i  -> Just $ Float64Lit $ fromIntegral i
      Int64Lit  i  -> Just $ Float64Lit $ fixUlp i $ fromIntegral i
      Word8Lit  i  -> Just $ Float64Lit $ fromIntegral i
      Word32Lit i  -> Just $ Float64Lit $ fromIntegral i
      Word64Lit i  -> Just $ Float64Lit $ fixUlp i $ fromIntegral i
      Float32Lit f -> Just $ Float64Lit $ float2Double f
      Float64Lit _ -> Just l
      PtrLit   _ _ -> Nothing
    where
      -- When casting an integer type to a floating-point type of lower precision
      -- (e.g., int32 to float32), GHC between 7.8.3 and 9.2.2 (exclusive) rounds
      -- toward zero, instead of rounding to nearest even like everybody else.
      -- See https://gitlab.haskell.org/ghc/ghc/-/issues/17231.
      --
      -- We patch this by manually checking the two adjacent floats to the
      -- candidate answer, and using one of those if the reverse cast is closer
      -- to the original input.
      --
      -- This rounds to nearest.  We round to nearest *even* by considering the
      -- candidates in decreasing order of the number of trailing zeros they
      -- exhibit when cast back to the original integer type.
      fixUlp :: forall a b w. (Num a, Integral a, FiniteBits a, RealFrac b, FloatingBits b w)
        => a -> b -> b
      fixUlp orig candidate = res where
        res = closest $ sortBy moreLowBits [candidate, candidatem1, candidatep1]
        candidatem1 = nextDown candidate
        candidatep1 = nextUp candidate
        closest = minimumBy (\ca cb -> err ca `compare` err cb)
        err cand = absdiff orig (round cand)
        absdiff a b = if a >= b then a - b else b - a
        moreLowBits a b =
          compare (0 - countTrailingZeros (round @b @a a))
                  (0 - countTrailingZeros (round @b @a b))
  
  -- === Helpers for function evaluation over fixed-width types ===
  
  applyIntBinOp :: (forall a. (Num a, Integral a) => a -> a -> a) -> LitVal -> LitVal -> LitVal
  applyIntBinOp f x y = case (x, y) of
    (Int64Lit  x', Int64Lit  y') -> Int64Lit  $ f x' y'
    (Int32Lit  x', Int32Lit  y') -> Int32Lit  $ f x' y'
    (Word8Lit  x', Word8Lit  y') -> Word8Lit  $ f x' y'
    (Word32Lit x', Word32Lit y') -> Word32Lit $ f x' y'
    (Word64Lit x', Word64Lit y') -> Word64Lit $ f x' y'
    _ -> error "Expected integer atoms"
  
  applyIntCmpOp :: (forall a. (Eq a, Ord a) => a -> a -> Bool) -> LitVal -> LitVal -> LitVal
  applyIntCmpOp f x y = boolLit case (x, y) of
    (Int64Lit  x', Int64Lit  y') -> f x' y'
    (Int32Lit  x', Int32Lit  y') -> f x' y'
    (Word8Lit  x', Word8Lit  y') -> f x' y'
    (Word32Lit x', Word32Lit y') -> f x' y'
    (Word64Lit x', Word64Lit y') -> f x' y'
    _ -> error "Expected integer atoms"
  
  applyFloatBinOp :: (forall a. (Num a, Fractional a) => a -> a -> a) -> LitVal -> LitVal -> LitVal
  applyFloatBinOp f x y = case (x, y) of
    (Float64Lit x', Float64Lit y') -> Float64Lit $ f x' y'
    (Float32Lit x', Float32Lit y') -> Float32Lit $ f x' y'
    _ -> error "Expected float atoms"
  
  applyFloatCmpOp :: (forall a. (Eq a, Ord a) => a -> a -> Bool) -> LitVal -> LitVal -> LitVal
  applyFloatCmpOp f x y = boolLit case (x, y) of
    (Float64Lit x', Float64Lit y') -> f x' y'
    (Float32Lit x', Float32Lit y') -> f x' y'
    _ -> error "Expected float atoms"
  
  boolLit :: Bool -> LitVal
  boolLit x = Word8Lit $ fromIntegral $ fromEnum x
  
  cmp :: Ord a => CmpOp -> a -> a -> Bool
  cmp = \case
    Less         -> (<)
    Greater      -> (>)
    Equal        -> (==)
    LessEqual    -> (<=)
    GreaterEqual -> (>=)
  