0 | module TTImp.ProcessFnOpt

  2 | import Core.Context.Log

  3 | import Core.Env

  4 | import Core.Normalise

  5 | import Core.Value

  7 | import TTImp.TTImp

  9 | import Libraries.Data.NameMap

 10 | import Libraries.Data.NatSet

 12 | getRetTy : Defs -> ClosedNF -> Core Name

 13 | getRetTy defs (NBind fc _ (Pi {}) sc)

 14 |     = getRetTy defs !(sc defs (toClosure defaultOpts Env.empty (Erased fc Placeholder)))

 15 | getRetTy defs (NTCon _ n _ _) = pure n

 16 | getRetTy defs ty

 17 |     = throw (GenericMsg (getLoc ty)

 18 |              "Can only add hints for concrete return types")

 21 | throwIf : {auto c : Ref Ctxt Defs} -> FC -> Bool -> String -> Core ()

 22 | throwIf fc cond msg = when cond $ throw (GenericMsg fc msg)

 25 | throwIfHasFlag : {auto c : Ref Ctxt Defs} -> FC -> Name -> DefFlag -> String -> Core ()

 26 | throwIfHasFlag fc ndef fl msg = throwIf fc !(hasFlag fc ndef fl) msg

 29 | throwIfHasTotality : {auto c : Ref Ctxt Defs} -> FC -> Name -> String -> Core ()

 30 | throwIfHasTotality fc ndef msg = throwIf fc !(hasSetTotal fc ndef) msg

 33 | processFnOpt : {auto c : Ref Ctxt Defs} ->

 34 |                FC -> Bool -> -- ^ top level name?

 35 |                Name -> FnOpt -> Core ()

 36 | processFnOpt fc _ ndef Unsafe

 37 |     = do setIsEscapeHatch fc ndef

 38 | processFnOpt fc _ ndef Inline

 39 |     = do throwIfHasFlag fc ndef NoInline "%noinline and %inline are mutually exclusive"

 40 |          setFlag fc ndef Inline

 41 | processFnOpt fc _ ndef NoInline

 42 |     = do throwIfHasFlag fc ndef Inline "%inline and %noinline are mutually exclusive"

 43 |          setFlag fc ndef NoInline

 44 | processFnOpt fc _ ndef Deprecate

 45 |     = setFlag fc ndef Deprecate

 46 | processFnOpt fc _ ndef TCInline

 47 |     = setFlag fc ndef TCInline

 48 | processFnOpt fc True ndef (Hint d)

 49 |     = do defs <- get Ctxt

 50 |          Just ty <- lookupTyExact ndef (gamma defs)

 51 |               | Nothing => undefinedName fc ndef

 52 |          target <- getRetTy defs !(nf defs Env.empty ty)

 53 |          addHintFor fc target ndef d False

 54 | processFnOpt fc _ ndef (Hint d)

 55 |     = do logC "elab" 5 $ do pure $ "Adding local hint " ++ show !(toFullNames ndef)

 56 |          addLocalHint ndef

 57 | processFnOpt fc True ndef (GlobalHint a)

 58 |     = addGlobalHint ndef a

 59 | processFnOpt fc _ ndef (GlobalHint True)

 60 |     = throw (GenericMsg fc "%defaulthint is not valid in local definitions")

 61 | processFnOpt fc _ ndef (GlobalHint False)

 62 |     = throw (GenericMsg fc "%globalhint is not valid in local definitions")

 63 | processFnOpt fc _ ndef ExternFn

 64 |     = setFlag fc ndef Inline -- if externally defined, inline when compiling

 65 | processFnOpt fc _ ndef (ForeignFn _)

 66 |     = setFlag fc ndef Inline -- if externally defined, inline when compiling

 67 | processFnOpt fc _ ndef (ForeignExport _)

 68 |     = pure ()

 69 | processFnOpt fc _ ndef Invertible

 70 |     = setFlag fc ndef Invertible

 71 | processFnOpt fc _ ndef (Totality tot)

 72 |     = do throwIfHasTotality fc ndef "Multiple totality modifiers"

 73 |          setFlag fc ndef (SetTotal tot)

 74 | processFnOpt fc _ ndef Macro

 75 |     = setFlag fc ndef Macro

 76 | processFnOpt fc _ ndef (SpecArgs ns)

 77 |     = do defs <- get Ctxt

 78 |          Just gdef <- lookupCtxtExact ndef (gamma defs)

 79 |               | Nothing => undefinedName fc ndef

 80 |          nty <- nf defs Env.empty (type gdef)

 81 |          ps <- getNamePos 0 nty

 82 |          ddeps <- collectDDeps nty

 83 |          specs <- collectSpec NatSet.empty ddeps ps nty

 84 |          ignore $ addDef ndef ({ specArgs := specs } gdef)

 85 |   where

 86 |     insertDeps : NatSet -> List (Name, Nat) -> List Name -> NatSet

 87 |     insertDeps acc ps [] = acc

 88 |     insertDeps acc ps (n :: ns)

 89 |         = case lookup n ps of

 90 |                Nothing => insertDeps acc ps ns

 91 |                Just pos => if pos `elem` acc

 92 |                               then insertDeps acc ps ns

 93 |                               else insertDeps (NatSet.insert pos acc) ps ns

 95 |     -- Collect the argument names which the dynamic args depend on

 96 |     collectDDeps : ClosedNF -> Core (List Name)

 97 |     collectDDeps (NBind tfc x (Pi _ _ _ nty) sc)

 98 |         = do defs <- get Ctxt

 99 |              empty <- clearDefs defs

100 |              sc' <- sc defs (toClosure defaultOpts Env.empty (Ref tfc Bound x))

101 |              if x `elem` ns

102 |                 then collectDDeps sc'

103 |                 else do aty <- quote empty Env.empty nty

104 |                         -- Get names depended on by nty

105 |                         let deps = keys (getRefs (UN Underscore) aty)

106 |                         rest <- collectDDeps sc'

107 |                         pure (rest ++ deps)

108 |     collectDDeps _ = pure []

110 |     -- Return names the type depends on, and whether it's a parameter

111 |     mutual

112 |       getDepsArgs : Bool -> List ClosedNF -> NameMap Bool ->

113 |                     Core (NameMap Bool)

114 |       getDepsArgs inparam [] ns = pure ns

115 |       getDepsArgs inparam (a :: as) ns

116 |           = do ns' <- getDeps inparam a ns

117 |                getDepsArgs inparam as ns'

119 |       getDeps : Bool -> ClosedNF -> NameMap Bool ->

120 |                 Core (NameMap Bool)

121 |       getDeps inparam (NBind _ x (Pi _ _ _ pty) sc) ns

122 |           = do defs <- get Ctxt

123 |                ns' <- getDeps inparam !(evalClosure defs pty) ns

124 |                sc' <- sc defs (toClosure defaultOpts Env.empty (Erased fc Placeholder))

125 |                getDeps inparam sc' ns'

126 |       getDeps inparam (NBind _ x b sc) ns

127 |           = do defs <- get Ctxt

128 |                ns' <- getDeps False !(evalClosure defs (binderType b)) ns

129 |                sc' <- sc defs (toClosure defaultOpts Env.empty (Erased fc Placeholder))

130 |                getDeps False sc' ns

131 |       getDeps inparam (NApp _ (NRef Bound n) args) ns

132 |           = do defs <- get Ctxt

133 |                ns' <- getDepsArgs False !(traverse (evalClosure defs . snd) args) ns

134 |                pure (insert n inparam ns')

135 |       getDeps inparam (NDCon _ n t a args) ns

136 |           = do defs <- get Ctxt

137 |                getDepsArgs False !(traverse (evalClosure defs . snd) args) ns

138 |       getDeps inparam (NTCon _ n a args) ns

139 |           = do defs <- get Ctxt

140 |                params <- case !(lookupDefExact n (gamma defs)) of

141 |                               Just (TCon _ ps _ _ _ _ _) => pure ps

142 |                               _ => pure NatSet.empty

143 |                let (ps, ds) = NatSet.partition params (map snd args)

144 |                ns' <- getDepsArgs True !(traverse (evalClosure defs) ps) ns

145 |                getDepsArgs False !(traverse (evalClosure defs) ds) ns'

146 |       getDeps inparam (NDelayed _ _ t) ns = getDeps inparam t ns

147 |       getDeps inparams nf ns = pure ns

149 |     -- If the name of an argument is in the list of specialisable arguments,

150 |     -- record the position. Also record the position of anything the argument

151 |     -- depends on which is only dependend on by declared static arguments.

152 |     collectSpec : NatSet    -> -- specialisable so far

153 |                   List Name -> -- things depended on by dynamic args

154 |                                -- We're assuming it's a short list, so just use

155 |                                -- List and don't worry about duplicates.

156 |                   List (Name, Nat) -> ClosedNF -> Core NatSet

157 |     collectSpec acc ddeps ps (NBind tfc x (Pi _ _ _ nty) sc)

158 |         = do defs <- get Ctxt

159 |              empty <- clearDefs defs

160 |              sc' <- sc defs (toClosure defaultOpts Env.empty (Ref tfc Bound x))

161 |              if x `elem` ns

162 |                 then do deps <- getDeps True !(evalClosure defs nty) NameMap.empty

163 |                         -- Get names depended on by nty

164 |                         -- Keep the ones which are either:

165 |                         --  * parameters

166 |                         --  * not depended on by a dynamic argument (the ddeps)

167 |                         let rs = filter (\x => snd x ||

168 |                                                not (fst x `elem` ddeps))

169 |                                         (toList deps)

170 |                         let acc' = insertDeps acc ps (x :: map fst rs)

171 |                         collectSpec acc' ddeps ps sc'

172 |                 else collectSpec acc ddeps ps sc'

173 |     collectSpec acc ddeps ps _ = pure acc

175 |     getNamePos : Nat -> ClosedNF -> Core (List (Name, Nat))

176 |     getNamePos i (NBind tfc x (Pi {}) sc)

177 |         = do defs <- get Ctxt

178 |              ns' <- getNamePos (1 + i) !(sc defs (toClosure defaultOpts Env.empty (Erased tfc Placeholder)))

179 |              pure ((x, i) :: ns')

180 |     getNamePos _ _ = pure []