content: module Futhark.Analysis.AlgSimplify
    ( Prod (..),
      SofP,
      simplify0,
      simplify,
      simplify',
      simplifySofP,
      simplifySofP',
      sumOfProducts,
      sumToExp,
      prodToExp,
      add,
      sub,
      negate,
      isMultipleOf,
      maybeDivide,
      removeLessThans,
      lessThanish,
      compareComplexity,
    )
  where
  
  import Data.Bits (xor)
  import Data.Function ((&))
  import Data.List (findIndex, intersect, partition, sort, (\\))
  import Data.Maybe (mapMaybe)
  import Futhark.Analysis.PrimExp
  import Futhark.Analysis.PrimExp.Convert
  import Futhark.IR.Prop.Names
  import Futhark.IR.Syntax.Core (SubExp (..), VName)
  import Futhark.Util
  import Futhark.Util.Pretty
  import Prelude hiding (negate)
  
  type Exp = PrimExp VName
  
  type TExp = TPrimExp Int64 VName
  
  data Prod = Prod
    { negated :: Bool,
      atoms :: [Exp]
    }
    deriving (Show, Eq, Ord)
  
  type SofP = [Prod]
  
  sumOfProducts :: Exp -> SofP
  sumOfProducts = map sortProduct . sumOfProducts'
  
  sortProduct :: Prod -> Prod
  sortProduct (Prod n as) = Prod n $ sort as
  
  sumOfProducts' :: Exp -> SofP
  sumOfProducts' (BinOpExp (Add Int64 _) e1 e2) =
    sumOfProducts' e1 <> sumOfProducts' e2
  sumOfProducts' (BinOpExp (Sub Int64 _) (ValueExp (IntValue (Int64Value 0))) e) =
    map negate $ sumOfProducts' e
  sumOfProducts' (BinOpExp (Sub Int64 _) e1 e2) =
    sumOfProducts' e1 <> map negate (sumOfProducts' e2)
  sumOfProducts' (BinOpExp (Mul Int64 _) e1 e2) =
    sumOfProducts' e1 `mult` sumOfProducts' e2
  sumOfProducts' (ValueExp (IntValue (Int64Value i))) =
    [Prod (i < 0) [ValueExp $ IntValue $ Int64Value $ abs i]]
  sumOfProducts' e = [Prod False [e]]
  
  mult :: SofP -> SofP -> SofP
  mult xs ys = [Prod (b `xor` b') (x <> y) | Prod b x <- xs, Prod b' y <- ys]
  
  negate :: Prod -> Prod
  negate p = p {negated = not $ negated p}
  
  sumToExp :: SofP -> Exp
  sumToExp [] = val 0
  sumToExp [x] = prodToExp x
  sumToExp (x : xs) =
    foldl (BinOpExp $ Add Int64 OverflowUndef) (prodToExp x) $
      map prodToExp xs
  
  prodToExp :: Prod -> Exp
  prodToExp (Prod _ []) = val 1
  prodToExp (Prod True [ValueExp (IntValue (Int64Value i))]) = ValueExp $ IntValue $ Int64Value (-i)
  prodToExp (Prod True as) =
    foldl (BinOpExp $ Mul Int64 OverflowUndef) (val (-1)) as
  prodToExp (Prod False (a : as)) =
    foldl (BinOpExp $ Mul Int64 OverflowUndef) a as
  
  simplifySofP :: SofP -> SofP
  simplifySofP =
    -- TODO: Maybe 'constFoldValueExps' is not necessary after adding scaleConsts
    fixPoint (mapMaybe (applyZero . removeOnes) . scaleConsts . constFoldValueExps . removeNegations)
  
  simplifySofP' :: SofP -> SofP
  simplifySofP' = fixPoint (mapMaybe (applyZero . removeOnes) . scaleConsts . removeNegations)
  
  simplify0 :: Exp -> SofP
  simplify0 = simplifySofP . sumOfProducts
  
  simplify :: Exp -> Exp
  simplify = constFoldPrimExp . sumToExp . simplify0
  
  simplify' :: TExp -> TExp
  simplify' = TPrimExp . simplify . untyped
  
  applyZero :: Prod -> Maybe Prod
  applyZero p@(Prod _ as)
    | val 0 `elem` as = Nothing
    | otherwise = Just p
  
  removeOnes :: Prod -> Prod
  removeOnes (Prod neg as) =
    let as' = filter (/= val 1) as
     in Prod neg $ if null as' then [ValueExp $ IntValue $ Int64Value 1] else as'
  
  removeNegations :: SofP -> SofP
  removeNegations [] = []
  removeNegations (t : ts) =
    case break (== negate t) ts of
      (start, _ : rest) -> removeNegations $ start <> rest
      _ -> t : removeNegations ts
  
  constFoldValueExps :: SofP -> SofP
  constFoldValueExps prods =
    let (value_exps, others) = partition (all isPrimValue . atoms) prods
        value_exps' = sumOfProducts $ constFoldPrimExp $ sumToExp value_exps
     in value_exps' <> others
  
  intFromExp :: Exp -> Maybe Int64
  intFromExp (ValueExp (IntValue x)) = Just $ valueIntegral x
  intFromExp _ = Nothing
  
  -- | Given @-[2, x]@ returns @(-2, [x])@
  prodToScale :: Prod -> (Int64, [Exp])
  prodToScale (Prod b exps) =
    let (scalars, exps') = partitionMaybe intFromExp exps
     in if b
          then (-(product scalars), exps')
          else (product scalars, exps')
  
  -- | Given @(-2, [x])@ returns @-[1, 2, x]@
  scaleToProd :: (Int64, [Exp]) -> Prod
  scaleToProd (i, exps) =
    Prod (i < 0) $ ValueExp (IntValue $ Int64Value $ abs i) : exps
  
  -- | Given @[[2, x], -[x]]@ returns @[[x]]@
  scaleConsts :: SofP -> SofP
  scaleConsts =
    helper [] . map prodToScale
    where
      helper :: [Prod] -> [(Int64, [Exp])] -> [Prod]
      helper acc [] = reverse acc
      helper acc ((scale, exps) : rest) =
        case flip focusNth rest =<< findIndex ((==) exps . snd) rest of
          Nothing -> helper (scaleToProd (scale, exps) : acc) rest
          Just (before, (scale', _), after) ->
            helper acc $ (scale + scale', exps) : (before <> after)
  
  isPrimValue :: Exp -> Bool
  isPrimValue (ValueExp _) = True
  isPrimValue _ = False
  
  val :: Int64 -> Exp
  val = ValueExp . IntValue . Int64Value
  
  add :: SofP -> SofP -> SofP
  add ps1 ps2 = simplifySofP $ ps1 <> ps2
  
  sub :: SofP -> SofP -> SofP
  sub ps1 ps2 = add ps1 $ map negate ps2
  
  isMultipleOf :: Prod -> [Exp] -> Bool
  isMultipleOf (Prod _ as) term =
    let quotient = as \\ term
     in sort (quotient <> term) == sort as
  
  maybeDivide :: Prod -> Prod -> Maybe Prod
  maybeDivide dividend divisor
    | Prod dividend_b dividend_factors <- dividend,
      Prod divisor_b divisor_factors <- divisor,
      quotient <- dividend_factors \\ divisor_factors,
      sort (quotient <> divisor_factors) == sort dividend_factors =
        Just $ Prod (dividend_b `xor` divisor_b) quotient
    | (dividend_scale, dividend_rest) <- prodToScale dividend,
      (divisor_scale, divisor_rest) <- prodToScale divisor,
      dividend_scale `mod` divisor_scale == 0,
      null $ divisor_rest \\ dividend_rest =
        Just $
          Prod
            (signum (dividend_scale `div` divisor_scale) < 0)
            ( ValueExp (IntValue $ Int64Value $ dividend_scale `div` divisor_scale)
                : (dividend_rest \\ divisor_rest)
            )
    | otherwise = Nothing
  
  -- | Given a list of 'Names' that we know are non-negative (>= 0), determine
  -- whether we can say for sure that the given 'AlgSimplify.SofP' is
  -- non-negative. Conservatively returns 'False' if there is any doubt.
  --
  -- TODO: We need to expand this to be able to handle cases such as @i*n + g < (i
  -- + 1) * n@, if it is known that @g < n@, eg. from a 'SegSpace' or a loop form.
  nonNegativeish :: Names -> SofP -> Bool
  nonNegativeish non_negatives = all (nonNegativeishProd non_negatives)
  
  nonNegativeishProd :: Names -> Prod -> Bool
  nonNegativeishProd _ (Prod True _) = False
  nonNegativeishProd non_negatives (Prod False as) =
    all (nonNegativeishExp non_negatives) as
  
  nonNegativeishExp :: Names -> PrimExp VName -> Bool
  nonNegativeishExp _ (ValueExp v) = not $ negativeIsh v
  nonNegativeishExp non_negatives (LeafExp vname _) = vname `nameIn` non_negatives
  nonNegativeishExp _ _ = False
  
  -- | Is e1 symbolically less than or equal to e2?
  lessThanOrEqualish :: [(VName, PrimExp VName)] -> Names -> TPrimExp Int64 VName -> TPrimExp Int64 VName -> Bool
  lessThanOrEqualish less_thans0 non_negatives e1 e2 =
    case e2 - e1 & untyped & simplify0 of
      [] -> True
      simplified ->
        nonNegativeish non_negatives $
          fixPoint (`removeLessThans` less_thans) simplified
    where
      less_thans =
        concatMap
          (\(i, bound) -> [(Var i, bound), (Constant $ IntValue $ Int64Value 0, bound)])
          less_thans0
  
  lessThanish :: [(VName, PrimExp VName)] -> Names -> TPrimExp Int64 VName -> TPrimExp Int64 VName -> Bool
  lessThanish less_thans non_negatives e1 =
    lessThanOrEqualish less_thans non_negatives (e1 + 1)
  
  removeLessThans :: SofP -> [(SubExp, PrimExp VName)] -> SofP
  removeLessThans =
    foldl
      ( \sofp (i, bound) ->
          let to_remove =
                simplifySofP $
                  Prod True [primExpFromSubExp (IntType Int64) i]
                    : simplify0 bound
           in case to_remove `intersect` sofp of
                to_remove' | to_remove' == to_remove -> sofp \\ to_remove
                _ -> sofp
      )
  
  compareComplexity :: SofP -> SofP -> Ordering
  compareComplexity xs0 ys0 =
    case length xs0 `compare` length ys0 of
      EQ -> helper xs0 ys0
      c -> c
    where
      helper [] [] = EQ
      helper [] _ = LT
      helper _ [] = GT
      helper (px : xs) (py : ys) =
        case (prodToScale px, prodToScale py) of
          ((ix, []), (iy, [])) -> case ix `compare` iy of
            EQ -> helper xs ys
            c -> c
          ((_, []), (_, _)) -> LT
          ((_, _), (_, [])) -> GT
          ((_, x), (_, y)) -> case length x `compare` length y of
            EQ -> helper xs ys
            c -> c
  