{-|
Module      : Idris.DataOpts
Description : Optimisations for Idris code i.e. Forcing, detagging and collapsing.

License     : BSD3
Maintainer  : The Idris Community.
-}
{-# LANGUAGE FlexibleInstances, PatternGuards #-}

module Idris.DataOpts(applyOpts) where

import Idris.AbsSyntax
import Idris.Core.TT

class Optimisable term where
    applyOpts :: term -> Idris term

instance (Optimisable a, Optimisable b) => Optimisable (a, b) where
    applyOpts :: (a, b) -> Idris (a, b)
applyOpts (a
x, b
y) = (,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall term. Optimisable term => term -> Idris term
applyOpts a
x forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall term. Optimisable term => term -> Idris term
applyOpts b
y

instance (Optimisable a, Optimisable b) => Optimisable (vs, a, b) where
    applyOpts :: (vs, a, b) -> Idris (vs, a, b)
applyOpts (vs
v, a
x, b
y) = (,,) vs
v forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall term. Optimisable term => term -> Idris term
applyOpts a
x forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall term. Optimisable term => term -> Idris term
applyOpts b
y

instance Optimisable a => Optimisable [a] where
    applyOpts :: [a] -> Idris [a]
applyOpts = forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall term. Optimisable term => term -> Idris term
applyOpts

instance Optimisable a => Optimisable (Either a (a, a)) where
    applyOpts :: Either a (a, a) -> Idris (Either a (a, a))
applyOpts (Left  a
t) = forall a b. a -> Either a b
Left  forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall term. Optimisable term => term -> Idris term
applyOpts a
t
    applyOpts (Right (a, a)
t) = forall a b. b -> Either a b
Right forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall term. Optimisable term => term -> Idris term
applyOpts (a, a)
t

-- Raw is for compile time optimisation (before type checking)
-- Term is for run time optimisation (after type checking, collapsing allowed)
-- Compile time: no collapsing

instance Optimisable Raw where
    applyOpts :: Raw -> Idris Raw
applyOpts t :: Raw
t@(RApp Raw
f Raw
a)
        | (Var Name
n, [Raw]
args) <- Raw -> (Raw, [Raw])
raw_unapply Raw
t -- MAGIC HERE
            = Raw -> [Raw] -> Raw
raw_apply (Name -> Raw
Var Name
n) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall term. Optimisable term => term -> Idris term
applyOpts [Raw]
args
        | Bool
otherwise = Raw -> Raw -> Raw
RApp forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall term. Optimisable term => term -> Idris term
applyOpts Raw
f forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall term. Optimisable term => term -> Idris term
applyOpts Raw
a

    applyOpts (RBind Name
n Binder Raw
b Raw
t) = Name -> Binder Raw -> Raw -> Raw
RBind Name
n forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall term. Optimisable term => term -> Idris term
applyOpts Binder Raw
b forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall term. Optimisable term => term -> Idris term
applyOpts Raw
t
    applyOpts Raw
t = forall (m :: * -> *) a. Monad m => a -> m a
return Raw
t

-- Erase types (makes ibc smaller, and we don't need them)
instance Optimisable (Binder (TT Name)) where
    applyOpts :: Binder (TT Name) -> Idris (Binder (TT Name))
applyOpts (Let RigCount
r TT Name
t TT Name
v) = forall b. RigCount -> b -> b -> Binder b
Let RigCount
r forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a. Monad m => a -> m a
return forall n. TT n
Erased forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall term. Optimisable term => term -> Idris term
applyOpts TT Name
v
    applyOpts Binder (TT Name)
b = forall (m :: * -> *) a. Monad m => a -> m a
return (Binder (TT Name)
b { binderTy :: TT Name
binderTy = forall n. TT n
Erased })

instance Optimisable (Binder Raw) where
    applyOpts :: Binder Raw -> Idris (Binder Raw)
applyOpts Binder Raw
b = do Raw
t' <- forall term. Optimisable term => term -> Idris term
applyOpts (forall b. Binder b -> b
binderTy Binder Raw
b)
                     forall (m :: * -> *) a. Monad m => a -> m a
return (Binder Raw
b { binderTy :: Raw
binderTy = Raw
t' })

-- Run-time: do everything

prel :: [Text]
prel = [String -> Text
txt String
"Nat", String -> Text
txt String
"Prelude"]

instance Optimisable (TT Name) where
    applyOpts :: TT Name -> Idris (TT Name)
applyOpts (P NameType
_ (NS (UN Text
fn) [Text]
mod) TT Name
_)
       | Text
fn forall a. Eq a => a -> a -> Bool
== String -> Text
txt String
"plus" Bool -> Bool -> Bool
&& [Text]
mod forall a. Eq a => a -> a -> Bool
== [Text]
prel
        = forall (m :: * -> *) a. Monad m => a -> m a
return (forall n. NameType -> n -> TT n -> TT n
P NameType
Ref (String -> Name
sUN String
"prim__addBigInt") forall n. TT n
Erased)
    applyOpts (P NameType
_ (NS (UN Text
fn) [Text]
mod) TT Name
_)
       | Text
fn forall a. Eq a => a -> a -> Bool
== String -> Text
txt String
"mult" Bool -> Bool -> Bool
&& [Text]
mod forall a. Eq a => a -> a -> Bool
== [Text]
prel
        = forall (m :: * -> *) a. Monad m => a -> m a
return (forall n. NameType -> n -> TT n -> TT n
P NameType
Ref (String -> Name
sUN String
"prim__mulBigInt") forall n. TT n
Erased)
    applyOpts (P NameType
_ (NS (UN Text
fn) [Text]
mod) TT Name
_)
       | Text
fn forall a. Eq a => a -> a -> Bool
== String -> Text
txt String
"divNat" Bool -> Bool -> Bool
&& [Text]
mod forall a. Eq a => a -> a -> Bool
== [Text]
prel
        = forall (m :: * -> *) a. Monad m => a -> m a
return (forall n. NameType -> n -> TT n -> TT n
P NameType
Ref (String -> Name
sUN String
"prim__sdivBigInt") forall n. TT n
Erased)
    applyOpts (P NameType
_ (NS (UN Text
fn) [Text]
mod) TT Name
_)
       | Text
fn forall a. Eq a => a -> a -> Bool
== String -> Text
txt String
"modNat" Bool -> Bool -> Bool
&& [Text]
mod forall a. Eq a => a -> a -> Bool
== [Text]
prel
        = forall (m :: * -> *) a. Monad m => a -> m a
return (forall n. NameType -> n -> TT n -> TT n
P NameType
Ref (String -> Name
sUN String
"prim__sremBigInt") forall n. TT n
Erased)
    applyOpts (App AppStatus Name
_ (P NameType
_ (NS (UN Text
fn) [Text]
mod) TT Name
_) TT Name
x)
       | Text
fn forall a. Eq a => a -> a -> Bool
== String -> Text
txt String
"fromIntegerNat" Bool -> Bool -> Bool
&& [Text]
mod forall a. Eq a => a -> a -> Bool
== [Text]
prel
        = forall term. Optimisable term => term -> Idris term
applyOpts TT Name
x
    applyOpts (P NameType
_ (NS (UN Text
fn) [Text]
mod) TT Name
_)
       | Text
fn forall a. Eq a => a -> a -> Bool
== String -> Text
txt String
"fromIntegerNat" Bool -> Bool -> Bool
&& [Text]
mod forall a. Eq a => a -> a -> Bool
== [Text]
prel
        = forall (m :: * -> *) a. Monad m => a -> m a
return (forall n. AppStatus n -> TT n -> TT n -> TT n
App forall n. AppStatus n
Complete (forall n. NameType -> n -> TT n -> TT n
P NameType
Ref (Name -> [String] -> Name
sNS (String -> Name
sUN String
"id") [String
"Basics",String
"Prelude"]) forall n. TT n
Erased) forall n. TT n
Erased)
    applyOpts (P NameType
_ (NS (UN Text
fn) [Text]
mod) TT Name
_)
       | Text
fn forall a. Eq a => a -> a -> Bool
== String -> Text
txt String
"toIntegerNat" Bool -> Bool -> Bool
&& [Text]
mod forall a. Eq a => a -> a -> Bool
== [Text]
prel
         = forall (m :: * -> *) a. Monad m => a -> m a
return (forall n. AppStatus n -> TT n -> TT n -> TT n
App forall n. AppStatus n
Complete (forall n. NameType -> n -> TT n -> TT n
P NameType
Ref (Name -> [String] -> Name
sNS (String -> Name
sUN String
"id") [String
"Basics",String
"Prelude"]) forall n. TT n
Erased) forall n. TT n
Erased)
    applyOpts c :: TT Name
c@(P (DCon Int
t Int
arity Bool
uniq) Name
n TT Name
_)
        = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ Name -> Int -> Int -> Bool -> [TT Name] -> TT Name
applyDataOptRT Name
n Int
t Int
arity Bool
uniq []
    applyOpts t :: TT Name
t@(App AppStatus Name
s TT Name
f TT Name
a)
        | (c :: TT Name
c@(P (DCon Int
t Int
arity Bool
uniq) Name
n TT Name
_), [TT Name]
args) <- forall n. TT n -> (TT n, [TT n])
unApply TT Name
t
            = Name -> Int -> Int -> Bool -> [TT Name] -> TT Name
applyDataOptRT Name
n Int
t Int
arity Bool
uniq forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall term. Optimisable term => term -> Idris term
applyOpts [TT Name]
args
        | Bool
otherwise = forall n. AppStatus n -> TT n -> TT n -> TT n
App AppStatus Name
s forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall term. Optimisable term => term -> Idris term
applyOpts TT Name
f forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall term. Optimisable term => term -> Idris term
applyOpts TT Name
a
    applyOpts (Bind Name
n Binder (TT Name)
b TT Name
t) = forall n. n -> Binder (TT n) -> TT n -> TT n
Bind Name
n forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall term. Optimisable term => term -> Idris term
applyOpts Binder (TT Name)
b forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall term. Optimisable term => term -> Idris term
applyOpts TT Name
t
    applyOpts (Proj TT Name
t Int
i) = forall n. TT n -> Int -> TT n
Proj forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall term. Optimisable term => term -> Idris term
applyOpts TT Name
t forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
i
    applyOpts TT Name
t = forall (m :: * -> *) a. Monad m => a -> m a
return TT Name
t

-- | Need to saturate arguments first to ensure that optimisation happens uniformly
applyDataOptRT :: Name -> Int -> Int -> Bool -> [Term] -> Term
applyDataOptRT :: Name -> Int -> Int -> Bool -> [TT Name] -> TT Name
applyDataOptRT Name
n Int
tag Int
arity Bool
uniq [TT Name]
args
    | forall (t :: * -> *) a. Foldable t => t a -> Int
length [TT Name]
args forall a. Eq a => a -> a -> Bool
== Int
arity = Name -> [TT Name] -> TT Name
doOpts Name
n [TT Name]
args
    | Bool
otherwise = let extra :: [Name]
extra = Int -> [Name]
satArgs (Int
arity forall a. Num a => a -> a -> a
- forall (t :: * -> *) a. Foldable t => t a -> Int
length [TT Name]
args)
                      tm :: TT Name
tm = Name -> [TT Name] -> TT Name
doOpts Name
n (forall a b. (a -> b) -> [a] -> [b]
map (forall n. Int -> TT n -> TT n
weakenTm (forall (t :: * -> *) a. Foldable t => t a -> Int
length [Name]
extra)) [TT Name]
args forall a. [a] -> [a] -> [a]
++ forall a b. (a -> b) -> [a] -> [b]
map (\Name
n -> forall n. NameType -> n -> TT n -> TT n
P NameType
Bound Name
n forall n. TT n
Erased) [Name]
extra)
                  in forall {n}. Eq n => [n] -> TT n -> TT n
bind [Name]
extra TT Name
tm
  where
    satArgs :: Int -> [Name]
satArgs Int
n = forall a b. (a -> b) -> [a] -> [b]
map (\Int
i -> Int -> String -> Name
sMN Int
i String
"sat") [Int
1..Int
n]

    bind :: [n] -> TT n -> TT n
bind [] TT n
tm = TT n
tm
    bind (n
n:[n]
ns) TT n
tm = forall n. n -> Binder (TT n) -> TT n -> TT n
Bind n
n (forall b. RigCount -> b -> Binder b
Lam RigCount
RigW forall n. TT n
Erased) (forall n. Eq n => n -> TT n -> TT n
pToV n
n ([n] -> TT n -> TT n
bind [n]
ns TT n
tm))

    -- Nat special cases
    -- TODO: Would be nice if this was configurable in idris source!
    -- Issue #1597 https://github.com/idris-lang/Idris-dev/issues/1597
    doOpts :: Name -> [TT Name] -> TT Name
doOpts (NS (UN Text
z) [Text
nat, Text
prelude]) []
        | Text
z forall a. Eq a => a -> a -> Bool
== String -> Text
txt String
"Z" Bool -> Bool -> Bool
&& Text
nat forall a. Eq a => a -> a -> Bool
== String -> Text
txt String
"Nat" Bool -> Bool -> Bool
&& Text
prelude forall a. Eq a => a -> a -> Bool
== String -> Text
txt String
"Prelude"
          = forall n. Const -> TT n
Constant (Integer -> Const
BI Integer
0)
    doOpts (NS (UN Text
s) [Text
nat, Text
prelude]) [TT Name
k]
        | Text
s forall a. Eq a => a -> a -> Bool
== String -> Text
txt String
"S" Bool -> Bool -> Bool
&& Text
nat forall a. Eq a => a -> a -> Bool
== String -> Text
txt String
"Nat" Bool -> Bool -> Bool
&& Text
prelude forall a. Eq a => a -> a -> Bool
== String -> Text
txt String
"Prelude"
          = forall n. AppStatus n -> TT n -> TT n -> TT n
App forall n. AppStatus n
Complete (forall n. AppStatus n -> TT n -> TT n -> TT n
App forall n. AppStatus n
Complete (forall n. NameType -> n -> TT n -> TT n
P NameType
Ref (String -> Name
sUN String
"prim__addBigInt") forall n. TT n
Erased) TT Name
k) (forall n. Const -> TT n
Constant (Integer -> Const
BI Integer
1))

    doOpts Name
n [TT Name]
args = forall n. TT n -> [TT n] -> TT n
mkApp (forall n. NameType -> n -> TT n -> TT n
P (Int -> Int -> Bool -> NameType
DCon Int
tag Int
arity Bool
uniq) Name
n forall n. TT n
Erased) [TT Name]
args