    module Agda.Compiler.Treeless.Identity
      ( detectIdentityFunctions ) where

    import Prelude hiding ((!!))  -- don't use partial functions

    import Control.Applicative ( Alternative((<|>), empty) )
    import Data.Semigroup
    import Data.List.NonEmpty (NonEmpty(..))
    import qualified Data.List as List

    import Agda.Syntax.Treeless
    import Agda.TypeChecking.Monad

    import Agda.Utils.List

    import Agda.Utils.Impossible

    detectIdentityFunctions :: QName -> TTerm -> TCM TTerm
    detectIdentityFunctions :: QName -> TTerm -> TCM TTerm
    detectIdentityFunctions QName
    q TTerm
    t =
      case QName -> TTerm -> Maybe (Int, Int)
    isIdentity QName
    q TTerm
    t of
        Maybe (Int, Int)
    Nothing     -> TTerm -> TCM TTerm
    forall (m :: * -> *) a. Monad m => a -> m a
    return TTerm
    t
        Just (Int
    n, Int
    k) -> do
          Bool -> QName -> TCM ()
    markInline Bool
    True QName
    q
          Defn
    def <- Definition -> Defn
    theDef (Definition -> Defn) -> TCMT IO Definition -> TCMT IO Defn
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> QName -> TCMT IO Definition
    forall (m :: * -> *). HasConstInfo m => QName -> m Definition
    getConstInfo QName
    q
          TTerm -> TCM TTerm
    forall (m :: * -> *) a. Monad m => a -> m a
    return (TTerm -> TCM TTerm) -> TTerm -> TCM TTerm
    forall a b. (a -> b) -> a -> b
    $ Int -> TTerm -> TTerm
    mkTLam Int
    n (TTerm -> TTerm) -> TTerm -> TTerm
    forall a b. (a -> b) -> a -> b
    $ Int -> TTerm
    TVar Int
    k

    -- If isIdentity f t = Just (n, k) then
    -- f = t is equivalent to f = λ xn₋₁ .. x₀ → xk
    isIdentity :: QName -> TTerm -> Maybe (Int, Int)
    isIdentity :: QName -> TTerm -> Maybe (Int, Int)
    isIdentity QName
    q TTerm
    t =
      QName -> TTerm -> Maybe (Int, Int)
    trivialIdentity QName
    q TTerm
    t Maybe (Int, Int) -> Maybe (Int, Int) -> Maybe (Int, Int)
    forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
    <|> QName -> TTerm -> Maybe (Int, Int)
    recursiveIdentity QName
    q TTerm
    t

    -- Does the function recurse on an argument, rebuilding the same value again.
    recursiveIdentity :: QName -> TTerm -> Maybe (Int, Int)
    recursiveIdentity :: QName -> TTerm -> Maybe (Int, Int)
    recursiveIdentity QName
    q TTerm
    t =
      case TTerm
    b of
        TCase Int
    x CaseInfo
    _ (TError TError
    TUnreachable) [TAlt]
    bs
          | (TAlt -> Bool) -> [TAlt] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all (Int -> TAlt -> Bool
    identityBranch Int
    x) [TAlt]
    bs -> (Int, Int) -> Maybe (Int, Int)
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Int
    n, Int
    x)
        TTerm
    _ -> Maybe (Int, Int)
    forall (f :: * -> *) a. Alternative f => f a
    empty -- TODO: lets?
      where
        (Int
    n, TTerm
    b) = TTerm -> (Int, TTerm)
    tLamView TTerm
    t

        identityBranch :: Int -> TAlt -> Bool
    identityBranch Int
    _ TALit{}   = Bool
    False
        identityBranch Int
    _ TAGuard{} = Bool
    False
        identityBranch Int
    x (TACon QName
    c Int
    a TTerm
    b) =
          case TTerm
    b of
            TApp (TCon QName
    c') Args
    args -> QName
    c QName -> QName -> Bool
    forall a. Eq a => a -> a -> Bool
    == QName
    c' Bool -> Bool -> Bool
    && Int -> Args -> Bool
    identityArgs Int
    a Args
    args
            TVar Int
    y              -> Int
    y Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == Int
    x Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + Int
    a  -- from @-pattern recovery
            TTerm
    _                   -> Bool
    False -- TODO: nested cases
          where
            identityArgs :: Int -> Args -> Bool
    identityArgs Int
    a Args
    args =
              Args -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length Args
    args Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == Int
    a Bool -> Bool -> Bool
    && [Bool] -> Bool
    forall (t :: * -> *). Foldable t => t Bool -> Bool
    and ((TTerm -> Int -> Bool) -> Args -> [Int] -> [Bool]
    forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
    zipWith TTerm -> Int -> Bool
    match (Args -> Args
    forall a. [a] -> [a]
    reverse Args
    args) [Int
    0..])

            proj :: Int -> [a] -> a
    proj Int
    x [a]
    args = a -> [a] -> Int -> a
    forall a. a -> [a] -> Int -> a
    indexWithDefault a
    forall a. HasCallStack => a
    __IMPOSSIBLE__ ([a] -> [a]
    forall a. [a] -> [a]
    reverse [a]
    args) Int
    x

            match :: TTerm -> Int -> Bool
    match TTerm
    TErased              Int
    _  = Bool
    True
            match (TVar Int
    z)             Int
    y = Int
    z Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == Int
    y
            match (TApp (TDef QName
    f) Args
    args) Int
    y = QName
    f QName -> QName -> Bool
    forall a. Eq a => a -> a -> Bool
    == QName
    q Bool -> Bool -> Bool
    && Args -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length Args
    args Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == Int
    n Bool -> Bool -> Bool
    && TTerm -> Int -> Bool
    match (Int -> Args -> TTerm
    forall {a}. Int -> [a] -> a
    proj Int
    x Args
    args) Int
    y
            match TTerm
    _                    Int
    _ = Bool
    False

    data IdentityIn = IdIn [Int]

    notId :: IdentityIn
    notId :: IdentityIn
    notId = [Int] -> IdentityIn
    IdIn []

    instance Semigroup IdentityIn where
      IdIn [Int]
    xs <> :: IdentityIn -> IdentityIn -> IdentityIn
    <> IdIn [Int]
    ys = [Int] -> IdentityIn
    IdIn ([Int] -> IdentityIn) -> [Int] -> IdentityIn
    forall a b. (a -> b) -> a -> b
    $ [Int] -> [Int] -> [Int]
    forall a. Eq a => [a] -> [a] -> [a]
    List.intersect [Int]
    xs [Int]
    ys

    -- Does the function always return one of its arguments unchanged (possibly
    -- through recursive calls).
    trivialIdentity :: QName -> TTerm -> Maybe (Int, Int)
    trivialIdentity :: QName -> TTerm -> Maybe (Int, Int)
    trivialIdentity QName
    q TTerm
    t =
      case Int -> TTerm -> IdentityIn
    go Int
    0 TTerm
    b of
        IdIn [Int
    x]     -> (Int, Int) -> Maybe (Int, Int)
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Int
    n, Int
    x)
        IdIn []      -> Maybe (Int, Int)
    forall a. Maybe a
    Nothing
        IdIn (Int
    _:Int
    _:[Int]
    _) -> Maybe (Int, Int)
    forall a. Maybe a
    Nothing  -- only happens for empty functions (which will never be called)
      where
        (Int
    n, TTerm
    b) = TTerm -> (Int, TTerm)
    tLamView TTerm
    t

        go :: Int -> TTerm -> IdentityIn
        go :: Int -> TTerm -> IdentityIn
    go Int
    k TTerm
    t =
          case TTerm
    t of
            TVar Int
    x | Int
    x Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    >= Int
    k    -> [Int] -> IdentityIn
    IdIn [Int
    x Int -> Int -> Int
    forall a. Num a => a -> a -> a
    - Int
    k]
                   | Bool
    otherwise -> IdentityIn
    notId
            TLet TTerm
    _ TTerm
    b           -> Int -> TTerm -> IdentityIn
    go (Int
    k Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + Int
    1) TTerm
    b
            TCase Int
    _ CaseInfo
    _ TTerm
    d [TAlt]
    bs     -> NonEmpty IdentityIn -> IdentityIn
    forall a. Semigroup a => NonEmpty a -> a
    sconcat (Int -> TTerm -> IdentityIn
    go Int
    k TTerm
    d IdentityIn -> [IdentityIn] -> NonEmpty IdentityIn
    forall a. a -> [a] -> NonEmpty a
    :| (TAlt -> IdentityIn) -> [TAlt] -> [IdentityIn]
    forall a b. (a -> b) -> [a] -> [b]
    map (Int -> TAlt -> IdentityIn
    goAlt Int
    k) [TAlt]
    bs)
            TApp (TDef QName
    f) Args
    args
              | QName
    f QName -> QName -> Bool
    forall a. Eq a => a -> a -> Bool
    == QName
    q         -> [Int] -> IdentityIn
    IdIn [ Int
    y | (TVar Int
    x, Int
    y) <- Args -> [Int] -> [(TTerm, Int)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip (Args -> Args
    forall a. [a] -> [a]
    reverse Args
    args) [Int
    0..], Int
    y Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + Int
    k Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == Int
    x ]
            TCoerce TTerm
    v          -> Int -> TTerm -> IdentityIn
    go Int
    k TTerm
    v
            TApp{}             -> IdentityIn
    notId
            TLam{}             -> IdentityIn
    notId
            TLit{}             -> IdentityIn
    notId
            TDef{}             -> IdentityIn
    notId
            TCon{}             -> IdentityIn
    notId
            TPrim{}            -> IdentityIn
    notId
            TUnit{}            -> IdentityIn
    notId
            TSort{}            -> IdentityIn
    notId
            TErased{}          -> IdentityIn
    notId
            TError{}           -> IdentityIn
    notId

        goAlt :: Int -> TAlt -> IdentityIn
        goAlt :: Int -> TAlt -> IdentityIn
    goAlt Int
    k (TALit Literal
    _ TTerm
    b)   = Int -> TTerm -> IdentityIn
    go Int
    k TTerm
    b
        goAlt Int
    k (TAGuard TTerm
    _ TTerm
    b) = Int -> TTerm -> IdentityIn
    go Int
    k TTerm
    b
        goAlt Int
    k (TACon QName
    _ Int
    n TTerm
    b) = Int -> TTerm -> IdentityIn
    go (Int
    k Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + Int
    n) TTerm
    b