Jikka-5.1.0.0: src/Jikka/Core/Convert/ConvexHullTrick.hs
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
-- |
-- Module : Jikka.Core.Convert.ConvexHullTrick
-- Description : uses convex hull trick. / convex hull trick を使います。
-- Copyright : (c) Kimiyuki Onaka, 2021
-- License : Apache License 2.0
-- Maintainer : kimiyuki95@gmail.com
-- Stability : experimental
-- Portability : portable
--
-- \[
-- \newcommand\int{\mathbf{int}}
-- \newcommand\bool{\mathbf{bool}}
-- \newcommand\list{\mathbf{list}}
-- \]
module Jikka.Core.Convert.ConvexHullTrick
( run,
-- * internal rules
rule,
parseLinearFunctionBody,
parseLinearFunctionBody',
)
where
import Control.Monad.Trans.Maybe
import Jikka.Common.Alpha
import Jikka.Common.Error
import Jikka.Core.Language.ArithmeticalExpr
import Jikka.Core.Language.Beta
import Jikka.Core.Language.BuiltinPatterns
import Jikka.Core.Language.Expr
import Jikka.Core.Language.FreeVars
import Jikka.Core.Language.Lint
import Jikka.Core.Language.RewriteRules
import Jikka.Core.Language.Util
-- | This is something commutative because only one kind of @c@ is allowed.
plusPair :: (ArithmeticalExpr, ArithmeticalExpr) -> (ArithmeticalExpr, ArithmeticalExpr) -> Maybe (ArithmeticalExpr, ArithmeticalExpr)
plusPair (a1, c1) (a2, _) | isZeroArithmeticalExpr a2 = Just (a1, c1)
plusPair (a1, c1) (_, c2) | isZeroArithmeticalExpr c2 = Just (a1, c1)
plusPair (a1, _) (a2, c2) | isZeroArithmeticalExpr a1 = Just (a2, c2)
plusPair (_, c1) (a2, c2) | isZeroArithmeticalExpr c1 = Just (a2, c2)
plusPair (a1, c1) (a2, c2) =
let (k1, c1') = splitConstantFactorArithmeticalExpr c1
(k2, c2') = splitConstantFactorArithmeticalExpr c2
a1' = multArithmeticalExpr (integerArithmeticalExpr k1) a1
a2' = multArithmeticalExpr (integerArithmeticalExpr k2) a2
in if c1' == c2'
then Just (plusArithmeticalExpr a1' a2', c1')
else Nothing
sumPairs :: [(ArithmeticalExpr, ArithmeticalExpr)] -> Maybe (ArithmeticalExpr, ArithmeticalExpr)
sumPairs = foldr (\e1 e2 -> plusPair e1 =<< e2) (Just (integerArithmeticalExpr 1, integerArithmeticalExpr 0))
-- | `parseLinearFunctionBody'` parses the body of a linear function which can be decomposed to convex hull trick.
-- @parseLinearFunctionBody' f i j e@ finds a 4-tuple @a, b, c, d@ where @e = a(f[j], j) c(f[< i], i) + b(f[j], j) + d(f[< i], i)@.
--
-- TODO: What is the relation between @j@ and @k@?
parseLinearFunctionBody' :: VarName -> VarName -> VarName -> Expr -> Maybe (Expr, Expr, Expr, Expr)
parseLinearFunctionBody' f i j e = result <$> go e
where
result (a, c, b, d) =
let (k, a') = splitConstantFactorArithmeticalExpr a
c' = multArithmeticalExpr (integerArithmeticalExpr k) c
in (formatArithmeticalExpr a', formatArithmeticalExpr c', formatArithmeticalExpr b, formatArithmeticalExpr d)
go = \case
Negate' e -> do
(a, c, b, d) <- go e
return (a, negateArithmeticalExpr c, negateArithmeticalExpr b, negateArithmeticalExpr d)
Plus' e1 e2 -> do
(a1, c1, b1, d1) <- go e1
(a2, c2, b2, d2) <- go e2
(a, c) <- plusPair (a1, c1) (a2, c2)
return (a, c, plusArithmeticalExpr b1 b2, plusArithmeticalExpr d1 d2)
Minus' e1 e2 -> do
(a1, c1, b1, d1) <- go e1
(a2, c2, b2, d2) <- go e2
(a, c) <- plusPair (a1, c1) (negateArithmeticalExpr a2, c2)
return (a, c, minusArithmeticalExpr b1 b2, minusArithmeticalExpr d1 d2)
Mult' e1 e2 -> do
(a1, c1, b1, d1) <- go e1
(a2, c2, b2, d2) <- go e2
(a, c) <-
sumPairs
[ (multArithmeticalExpr a1 a2, multArithmeticalExpr c1 c2),
(multArithmeticalExpr b2 a1, c1),
(multArithmeticalExpr b1 a2, c2),
(a1, multArithmeticalExpr c1 d2),
(a2, multArithmeticalExpr c2 d1),
(b2, d1),
(b1, d2)
]
return (a, c, multArithmeticalExpr b1 b2, multArithmeticalExpr d1 d2)
e
| f `isUnusedVar` e && j `isUnusedVar` e ->
-- NOTE: Put constants to @d@ and simplify @a, b@
return (integerArithmeticalExpr 1, integerArithmeticalExpr 0, integerArithmeticalExpr 0, parseArithmeticalExpr e)
e
| f `isUnusedVar` e && i `isUnusedVar` e ->
return (integerArithmeticalExpr 1, integerArithmeticalExpr 0, parseArithmeticalExpr e, integerArithmeticalExpr 0)
e@(At' _ (Var f') index) | f' == f -> case unNPlusKPattern (parseArithmeticalExpr index) of
Just (i', k) | i' == i && k < 0 -> do
return (integerArithmeticalExpr 1, integerArithmeticalExpr 0, integerArithmeticalExpr 0, parseArithmeticalExpr e)
Just (j', 0) | j' == j -> do
return (integerArithmeticalExpr 1, integerArithmeticalExpr 0, parseArithmeticalExpr e, integerArithmeticalExpr 0)
_ -> Nothing
_ -> Nothing
parseLinearFunctionBody :: MonadAlpha m => VarName -> VarName -> Integer -> Expr -> m (Maybe (Expr, Expr, Expr, Expr, Expr, Maybe Expr))
parseLinearFunctionBody f i k = runMaybeT . go
where
goMin e j step size = case unNPlusKPattern (parseArithmeticalExpr size) of
Just (i', k') | i' == i && k' == k -> do
(a, b, c, d) <- hoistMaybe $ parseLinearFunctionBody' f i j step
-- raname @j@ to @i@
a <- lift $ substitute j (Var i) a
c <- lift $ substitute j (Var i) c
return (LitInt' 1, a, b, c, d, (`Minus'` d) <$> e)
_ -> hoistMaybe Nothing
goMax e j step size = do
(sign, a, b, c, d, e) <- goMin e j step size
return (Negate' sign, a, Negate' b, Negate' c, d, Negate' <$> e)
go = \case
Min1' _ (Map' _ _ (Lam j _ step) (Range1' size)) -> goMin Nothing j step size
Max1' _ (Map' _ _ (Lam j _ step) (Range1' size)) -> goMax Nothing j step size
Min1' _ (Cons' _ e (Map' _ _ (Lam j _ step) (Range1' size))) -> goMin (Just e) j step size
Max1' _ (Cons' _ e (Map' _ _ (Lam j _ step) (Range1' size))) -> goMax (Just e) j step size
Min1' _ (Snoc' _ (Map' _ _ (Lam j _ step) (Range1' size)) e) -> goMin (Just e) j step size
Max1' _ (Snoc' _ (Map' _ _ (Lam j _ step) (Range1' size)) e) -> goMax (Just e) j step size
Negate' e -> do
(sign, a, b, c, d, e) <- go e
return (Negate' sign, a, b, c, Negate' d, e)
Plus' e1 e2 | isConstantTimeExpr e2 -> do
(sign, a, b, c, d, e) <- go e1
return (sign, a, b, c, Plus' d e2, e)
Plus' e1 e2 | isConstantTimeExpr e1 -> do
(sign, a, b, c, d, e) <- go e2
return (sign, a, b, c, Plus' e1 d, e)
Minus' e1 e2 | isConstantTimeExpr e2 -> do
(sign, a, b, c, d, e) <- go e1
return (sign, a, b, c, Minus' d e2, e)
Minus' e1 e2 | isConstantTimeExpr e1 -> do
(sign, a, b, c, d, e) <- go e2
return (Negate' sign, a, b, c, Minus' e1 d, e)
Mult' e1 e2 | isConstantTimeExpr e2 -> do
(sign, a, b, c, d, e) <- go e1
return (Mult' sign e2, a, b, c, Mult' d e2, e)
Mult' e1 e2 | isConstantTimeExpr e1 -> do
(sign, a, b, c, d, e) <- go e2
return (Mult' e1 sign, a, b, c, Mult' e1 d, e)
_ -> hoistMaybe Nothing
getLength :: Expr -> Maybe Integer
getLength = \case
Nil' _ -> Just 0
Cons' _ _ xs -> succ <$> getLength xs
Snoc' _ xs _ -> succ <$> getLength xs
_ -> Nothing
rule :: (MonadAlpha m, MonadError Error m) => RewriteRule m
rule = RewriteRule $ \_ -> \case
-- build (fun f -> step(f)) base n
Build' IntTy (Lam f _ step) base n -> runMaybeT $ do
let ts = [ConvexHullTrickTy, ListTy IntTy]
i <- lift genVarName'
k <- hoistMaybe $ getLength base
step <- replaceLenF f i k step
-- step(f) = sign() * min (cons e(f, i) (map (fun j -> a(f, j) c(f, i) + b(f, j)) (range (i + k)))) + d(f, i)
(sign, a, c, b, d, e) <- MaybeT $ parseLinearFunctionBody f i k step
-- Update base when k = 0. If user's program has no bugs, it uses min(cons(x, xs)) when k = 0.
(base, n, k, c, d, e) <- case (e, k) of
(Just e, 0) -> do
e0 <- lift $ substitute i (LitInt' 0) e
d0 <- lift $ substitute i (LitInt' 0) d
let base' = Let f (ListTy IntTy) base $ Snoc' IntTy base (Plus' (Mult' sign e0) d0)
c <- lift $ substitute i (Plus' (Var i) (LitInt' 1)) c
d <- lift $ substitute i (Plus' (Var i) (LitInt' 1)) d
e <- lift $ substitute i (Plus' (Var i) (LitInt' 1)) e
return (base', Minus' n (LitInt' 1), k + 1, c, d, Just e)
_ -> return (base, n, k, c, d, e)
-- base' = (cht, base)
base' <- do
x <- lift genVarName'
f' <- lift $ genVarName f
i' <- lift $ genVarName i
a <- lift $ substitute f (Var f') a
b <- lift $ substitute f (Var f') b
a <- lift $ substitute i (Var i') a
b <- lift $ substitute i (Var i') b
-- cht for base[0], ..., base[k - 1]
let cht = Foldl' IntTy ConvexHullTrickTy (Lam2 x ConvexHullTrickTy i' IntTy (ConvexHullTrickInsert' (Var x) a b)) ConvexHullTrickInit' (Range1' (LitInt' k))
return $
Let f' (ListTy IntTy) base $
uncurryApp (Tuple' ts) [cht, Var f']
-- step' = fun (cht, f) i ->
-- let f' = setat f index(i) value(..)
-- in let cht' = update cht a(i) b(i)
-- in (cht', f')
step' <- do
x <- lift genVarName'
-- value(..) = (min e (min cht f[i + k] + c(i)))
let value = Plus' (Mult' sign (maybe id (\e -> Min2' IntTy e) e (ConvexHullTrickGetMin' (Proj' ts 0 (Var x)) c))) d
y <- lift genVarName'
f' <- lift $ genVarName f
a <- lift $ substitute f (Var f') a
b <- lift $ substitute f (Var f') b
a <- lift $ substitute i (Plus' (Var i) (LitInt' k)) a
b <- lift $ substitute i (Plus' (Var i) (LitInt' k)) b
return $
Lam2 x (TupleTy ts) i IntTy $
Let f (ListTy IntTy) (Proj' ts 1 (Var x)) $
Let f' (ListTy IntTy) (Snoc' IntTy (Var f) value) $
Let y ConvexHullTrickTy (ConvexHullTrickInsert' (Proj' ts 0 (Var x)) a b) $
uncurryApp (Tuple' ts) [Var y, Var f']
-- proj 1 (foldl step' base' (range (n - 1)))
return $ Proj' ts 1 (Foldl' IntTy (TupleTy ts) step' base' (Range1' n))
_ -> return Nothing
runProgram :: (MonadAlpha m, MonadError Error m) => Program -> m Program
runProgram = applyRewriteRuleProgram' rule
-- | `run` optimizes a DP which has the recurrence relation
-- \[
-- \mathrm{dp}(i) = \min a(j) x(i) + b(j) \lbrace \mid j \lt i \rbrace + c(i)
-- \] where only appropriate elements of \(\mathrm{dp}\) are used in \(a, x, b, c\).
run :: (MonadAlpha m, MonadError Error m) => Program -> m Program
run prog = wrapError' "Jikka.Core.Convert.ConvexHullTrick" $ do
precondition $ do
ensureWellTyped prog
prog <- runProgram prog
postcondition $ do
ensureWellTyped prog
return prog