git.iko.soy/Idris-dev
on 385dcbf help sign in
Idris-dev/src/Idris/DataOpts.hs
118 lines 4.6 kB history raw
1{-|
2Module : Idris.DataOpts
3Description : Optimisations for Idris code i.e. Forcing, detagging and collapsing.
4Copyright :
5License : BSD3
6Maintainer : The Idris Community.
7-}
8{-# LANGUAGE PatternGuards #-}
10module Idris.DataOpts(applyOpts) where
12import Idris.AbsSyntax
13import Idris.AbsSyntaxTree
14import Idris.Core.TT
16import Control.Applicative
17import Data.List
18import Data.Maybe
19import Debug.Trace
21class Optimisable term where
22 applyOpts :: term -> Idris term
24instance (Optimisable a, Optimisable b) => Optimisable (a, b) where
25 applyOpts (x, y) = (,) <$> applyOpts x <*> applyOpts y
27instance (Optimisable a, Optimisable b) => Optimisable (vs, a, b) where
28 applyOpts (v, x, y) = (,,) v <$> applyOpts x <*> applyOpts y
30instance Optimisable a => Optimisable [a] where
31 applyOpts = mapM applyOpts
33instance Optimisable a => Optimisable (Either a (a, a)) where
34 applyOpts (Left t) = Left <$> applyOpts t
35 applyOpts (Right t) = Right <$> applyOpts t
37-- Raw is for compile time optimisation (before type checking)
38-- Term is for run time optimisation (after type checking, collapsing allowed)
39-- Compile time: no collapsing
41instance Optimisable Raw where
42 applyOpts t@(RApp f a)
43 | (Var n, args) <- raw_unapply t -- MAGIC HERE
44 = raw_apply (Var n) <$> mapM applyOpts args
45 | otherwise = RApp <$> applyOpts f <*> applyOpts a
47 applyOpts (RBind n b t) = RBind n <$> applyOpts b <*> applyOpts t
48 applyOpts t = return t
50-- Erase types (makes ibc smaller, and we don't need them)
51instance Optimisable (Binder (TT Name)) where
52 applyOpts (Let t v) = Let <$> return Erased <*> applyOpts v
53 applyOpts b = return (b { binderTy = Erased })
55instance Optimisable (Binder Raw) where
56 applyOpts b = do t' <- applyOpts (binderTy b)
57 return (b { binderTy = t' })
59-- Run-time: do everything
61prel = [txt "Nat", txt "Prelude"]
63instance Optimisable (TT Name) where
64 applyOpts (P _ (NS (UN fn) mod) _)
65 | fn == txt "plus" && mod == prel
66 = return (P Ref (sUN "prim__addBigInt") Erased)
67 applyOpts (P _ (NS (UN fn) mod) _)
68 | fn == txt "mult" && mod == prel
69 = return (P Ref (sUN "prim__mulBigInt") Erased)
70 applyOpts (P _ (NS (UN fn) mod) _)
71 | fn == txt "divNat" && mod == prel
72 = return (P Ref (sUN "prim__sdivBigInt") Erased)
73 applyOpts (P _ (NS (UN fn) mod) _)
74 | fn == txt "modNat" && mod == prel
75 = return (P Ref (sUN "prim__sremBigInt") Erased)
76 applyOpts (App _ (P _ (NS (UN fn) mod) _) x)
77 | fn == txt "fromIntegerNat" && mod == prel
78 = applyOpts x
79 applyOpts (P _ (NS (UN fn) mod) _)
80 | fn == txt "fromIntegerNat" && mod == prel
81 = return (App Complete (P Ref (sNS (sUN "id") ["Basics","Prelude"]) Erased) Erased)
82 applyOpts (P _ (NS (UN fn) mod) _)
83 | fn == txt "toIntegerNat" && mod == prel
84 = return (App Complete (P Ref (sNS (sUN "id") ["Basics","Prelude"]) Erased) Erased)
85 applyOpts c@(P (DCon t arity uniq) n _)
86 = return $ applyDataOptRT n t arity uniq []
87 applyOpts t@(App s f a)
88 | (c@(P (DCon t arity uniq) n _), args) <- unApply t
89 = applyDataOptRT n t arity uniq <$> mapM applyOpts args
90 | otherwise = App s <$> applyOpts f <*> applyOpts a
91 applyOpts (Bind n b t) = Bind n <$> applyOpts b <*> applyOpts t
92 applyOpts (Proj t i) = Proj <$> applyOpts t <*> pure i
93 applyOpts t = return t
95-- | Need to saturate arguments first to ensure that optimisation happens uniformly
96applyDataOptRT :: Name -> Int -> Int -> Bool -> [Term] -> Term
97applyDataOptRT n tag arity uniq args
98 | length args == arity = doOpts n args
99 | otherwise = let extra = satArgs (arity - length args)
100 tm = doOpts n (map (weakenTm (length extra)) args ++ map (\n -> P Bound n Erased) extra)
101 in bind extra tm
102 where
103 satArgs n = map (\i -> sMN i "sat") [1..n]
105 bind [] tm = tm
106 bind (n:ns) tm = Bind n (Lam RigW Erased) (pToV n (bind ns tm))
108 -- Nat special cases
109 -- TODO: Would be nice if this was configurable in idris source!
110 -- Issue #1597 https://github.com/idris-lang/Idris-dev/issues/1597
111 doOpts (NS (UN z) [nat, prelude]) []
112 | z == txt "Z" && nat == txt "Nat" && prelude == txt "Prelude"
113 = Constant (BI 0)
114 doOpts (NS (UN s) [nat, prelude]) [k]
115 | s == txt "S" && nat == txt "Nat" && prelude == txt "Prelude"
116 = App Complete (App Complete (P Ref (sUN "prim__addBigInt") Erased) k) (Constant (BI 1))
118 doOpts n args = mkApp (P (DCon tag arity uniq) n Erased) args