{-|
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 (x, y) = (,) <$> applyOpts x <*> applyOpts y

instance (Optimisable a, Optimisable b) => Optimisable (vs, a, b) where
    applyOpts (v, x, y) = (,,) v <$> applyOpts x <*> applyOpts y

instance Optimisable a => Optimisable [a] where
    applyOpts = mapM applyOpts

instance Optimisable a => Optimisable (Either a (a, a)) where
    applyOpts (Left  t) = Left  <$> applyOpts t
    applyOpts (Right t) = Right <$> applyOpts 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 t@(RApp f a)
        | (Var n, args) <- raw_unapply t -- MAGIC HERE
            = raw_apply (Var n) <$> mapM applyOpts args
        | otherwise = RApp <$> applyOpts f <*> applyOpts a

    applyOpts (RBind n b t) = RBind n <$> applyOpts b <*> applyOpts t
    applyOpts t = return t

-- Erase types (makes ibc smaller, and we don't need them)
instance Optimisable (Binder (TT Name)) where
    applyOpts (Let r t v) = Let r <$> return Erased <*> applyOpts v
    applyOpts b = return (b { binderTy = Erased })

instance Optimisable (Binder Raw) where
    applyOpts b = do t' <- applyOpts (binderTy b)
                     return (b { binderTy = t' })

-- Run-time: do everything

prel = [txt "Nat", txt "Prelude"]

instance Optimisable (TT Name) where
    applyOpts (P _ (NS (UN fn) mod) _)
       | fn == txt "plus" && mod == prel
        = return (P Ref (sUN "prim__addBigInt") Erased)
    applyOpts (P _ (NS (UN fn) mod) _)
       | fn == txt "mult" && mod == prel
        = return (P Ref (sUN "prim__mulBigInt") Erased)
    applyOpts (P _ (NS (UN fn) mod) _)
       | fn == txt "divNat" && mod == prel
        = return (P Ref (sUN "prim__sdivBigInt") Erased)
    applyOpts (P _ (NS (UN fn) mod) _)
       | fn == txt "modNat" && mod == prel
        = return (P Ref (sUN "prim__sremBigInt") Erased)
    applyOpts (App _ (P _ (NS (UN fn) mod) _) x)
       | fn == txt "fromIntegerNat" && mod == prel
        = applyOpts x
    applyOpts (P _ (NS (UN fn) mod) _)
       | fn == txt "fromIntegerNat" && mod == prel
        = return (App Complete (P Ref (sNS (sUN "id") ["Basics","Prelude"]) Erased) Erased)
    applyOpts (P _ (NS (UN fn) mod) _)
       | fn == txt "toIntegerNat" && mod == prel
         = return (App Complete (P Ref (sNS (sUN "id") ["Basics","Prelude"]) Erased) Erased)
    applyOpts c@(P (DCon t arity uniq) n _)
        = return $ applyDataOptRT n t arity uniq []
    applyOpts t@(App s f a)
        | (c@(P (DCon t arity uniq) n _), args) <- unApply t
            = applyDataOptRT n t arity uniq <$> mapM applyOpts args
        | otherwise = App s <$> applyOpts f <*> applyOpts a
    applyOpts (Bind n b t) = Bind n <$> applyOpts b <*> applyOpts t
    applyOpts (Proj t i) = Proj <$> applyOpts t <*> pure i
    applyOpts t = return t

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

    bind [] tm = tm
    bind (n:ns) tm = Bind n (Lam RigW Erased) (pToV n (bind ns 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 (NS (UN z) [nat, prelude]) []
        | z == txt "Z" && nat == txt "Nat" && prelude == txt "Prelude"
          = Constant (BI 0)
    doOpts (NS (UN s) [nat, prelude]) [k]
        | s == txt "S" && nat == txt "Nat" && prelude == txt "Prelude"
          = App Complete (App Complete (P Ref (sUN "prim__addBigInt") Erased) k) (Constant (BI 1))

    doOpts n args = mkApp (P (DCon tag arity uniq) n Erased) args