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