0 | module Deriving.DepTyCheck.Util.Specialisation
  1 |
  2 | import public Control.Monad.Either
  3 |
  4 | import public Data.DPair
  5 | import public Data.List.Ex
  6 | import public Data.List.Map
  7 | import public Data.SortedMap
  8 | import public Data.SortedMap.Extra
  9 | import public Data.SortedSet
 10 |
 11 | import public Deriving.DepTyCheck.Gen.ForOneType.Interface
 12 |
 13 | import public Deriving.SpecialiseData
 14 | import public Language.Reflection.Unify
 15 |
 16 | import public Data.Hashable
 17 | import public Data.Hashable.Base
 18 |
 19 | %default total
 20 |
 21 | allQImpl : Monad m => NamesInfoInTypes => TTImp -> TTImp -> m TTImp
 22 | allQImpl (IPi {}) r = pure r
 23 | allQImpl (IApp {}) (IApp _ (Implicit {}) _) = pure `(?)
 24 | allQImpl (IApp {}) r@(IApp {}) = pure r
 25 | allQImpl (IApp {}) _ = pure `(?)
 26 | allQImpl v@(IVar _ n) _ =
 27 |   case lookupType n of
 28 |     Just _ => pure v
 29 |     Nothing => pure `(?)
 30 | allQImpl _ _ = pure `(?)
 31 |
 32 | ||| Replace every non-function sub-expression with a question mark
 33 | |||
 34 | ||| (x -> (y -> z) -> q) becomes (? -> (? -> ?) -> ?)
 35 | allQuestions : NamesInfoInTypes => TTImp -> TTImp
 36 | allQuestions t = runIdentity $ mapMTTImp' allQImpl t
 37 |
 38 | ||| An abstract "argument" of a generator
 39 | |||
 40 | ||| Consists of a type constructor's argument and a possible given value
 41 | record GenArg where
 42 |   constructor MkGenArg
 43 |   arg : Arg
 44 |   given : Maybe TTImp
 45 |
 46 | LogPosition GenArg where
 47 |   logPosition (MkGenArg a Nothing) = "\{fromMaybe "<unnamed arg>" a.name}"
 48 |   logPosition (MkGenArg a $ Just t) = "(\{fromMaybe "<unnamed arg>" a.name} := \{show t})"
 49 |
 50 | unGA : List GenArg -> (List Arg, List (Maybe TTImp))
 51 | unGA [] = ([], [])
 52 | unGA (x :: xs) = let (ys, zs) = unGA xs in (x.arg :: ys, x.given :: zs)
 53 |
 54 | (.isGenerated) : GenArg -> Bool
 55 | (.isGenerated) = isNothing . given
 56 |
 57 | (.isGiven) : GenArg -> Bool
 58 | (.isGiven) = isJust . given
 59 |
 60 | ||| Determine if the argument should be specialised or passed through
 61 | (.isPassthrough) : Elaboration m => GenArg -> m Bool
 62 | (.isPassthrough) (MkGenArg a Nothing) = pure True
 63 | (.isPassthrough) (MkGenArg a $ Just g) = do
 64 |   let True = snd (unPi a.type) == `(Type)
 65 |     | _ => pure True
 66 |   case g of
 67 |     IVar _ n => do
 68 |       nInfo <- getInfo n
 69 |       case nInfo of
 70 |         [] => pure True
 71 |         _ => pure False
 72 |     _ => pure False
 73 |
 74 | ||| Assemble a list of arguments and their given values from `callGen` inputs
 75 | |||
 76 | ||| The indices inside both given lists must be in ascending order
 77 | mkArgs :
 78 |   NamesInfoInTypes =>
 79 |   (sig : GenSignature) ->
 80 |   List (Fin sig.targetType.args.length, Arg) ->
 81 |   List (Fin sig.targetType.args.length, TTImp) ->
 82 |   List GenArg
 83 | mkArgs sig [] _ = []
 84 | mkArgs sig ((_, x) :: xs) [] = MkGenArg x Nothing :: mkArgs sig xs []
 85 | mkArgs sig ((i1, x) :: xs) g@((i2, y) :: ys) =
 86 |   if i1 == i2
 87 |     then MkGenArg x (Just y) :: mkArgs sig xs ys
 88 |     else MkGenArg x Nothing  :: mkArgs sig xs g
 89 |
 90 | singleArg : NamesInfoInTypes => Nat -> GenArg -> (TTImp, List GenArg)
 91 | singleArg n (MkGenArg a v) = do
 92 |   let n : Name = fromString "lam^\{show n}"
 93 |   (IVar EmptyFC n, [MkGenArg (MkArg a.count a.piInfo (Just n) $ allQuestions a.type) v])
 94 |
 95 | processArg : MonadLog m => NamesInfoInTypes => GenSignature -> Nat -> GenArg -> m (TTImp, List GenArg)
 96 |
 97 | processArgs' : MonadLog m => NamesInfoInTypes => GenSignature -> Nat -> List GenArg -> m (List AnyApp, List GenArg)
 98 | processArgs' sig k [] = pure ([], [])
 99 | processArgs' sig k (x :: xs) = do
100 |   (aT, l) <- assert_total $ processArg sig k x
101 |   (recAA, l') <- processArgs' sig (k + length l) xs
102 |   pure (appArg x.arg aT :: recAA, l ++ l')
103 |
104 | processArg sig argIdx ga with (ga.given)
105 |   processArg sig argIdx ga | Nothing =
106 |     logValue DetailedDebug "deptycheck.derive.specialisation" [sig, ga]
107 |       "No given value, passing through"
108 |       $ singleArg argIdx ga
109 |   processArg sig argIdx ga | Just x = do
110 |     let (appLhs, appTerms) = unAppAny x
111 |     let IVar _ tyName = appLhs
112 |       | IPrimVal _ (PrT _) =>
113 |         logValue DetailedDebug "deptycheck.derive.specialisation" [sig, ga]
114 |           "Given a primitive type invocation, specialising"
115 |           (x, [])
116 |       | _ =>
117 |         logValue DetailedDebug "deptycheck.derive.specialisation" [sig, ga]
118 |           "Given value head is not a variable, passing through"
119 |           $ singleArg argIdx ga
120 |     case lookupType tyName of
121 |       Just tyInfo => do
122 |         let (_ :: _) = appTerms
123 |           | [] =>
124 |             logValue DetailedDebug "deptycheck.derive.specialisation" [sig, ga]
125 |               "Given a type invocation w/o arguments, specialising"
126 |               (x, [])
127 |         let givens = map (uncurry MkGenArg) $ zip tyInfo.args $ popArgVals tyInfo.args (mkAllApps appTerms)
128 |         logPoint DetailedDebug "deptycheck.derive.specialisation" [sig, ga]
129 |           "Given a type invocation, traversing arguments: \{show $ map (fromMaybe "" . name . arg) givens}"
130 |         map (mapFst $ reAppAny appLhs) $ processArgs' sig argIdx $ takeWhile (.isGiven) givens
131 |       Nothing => do
132 |         if (snd (unPi ga.arg.type) == `(Type))
133 |           then
134 |             logValue DetailedDebug "deptycheck.derive.specialisation" [sig, ga]
135 |               "Given a non-global type expr, passing through"
136 |               $ singleArg argIdx ga
137 |           else
138 |             logValue DetailedDebug "deptycheck.derive.specialisation" [sig, ga]
139 |               "Given a non-type expr, passing through"
140 |               $ singleArg argIdx ga
141 |
142 | processArgs :
143 |   MonadLog m =>
144 |   NamesInfoInTypes =>
145 |   (sig : GenSignature) ->
146 |   List GenArg ->
147 |   m (TTImp, List Arg, List $ Maybe TTImp)
148 | processArgs sig ga = bimap (reAppAny $ IVar EmptyFC sig.targetType.name) unGA <$> processArgs' sig 0 ga
149 |
150 | ||| Given a set of given argument indices, convert a list of their values into a vector that can be fed to `callGen`
151 | |||
152 | ||| The values should be listed for indices in ascending order
153 | ||| (i.e. how these indices would be sorted if we called `toList` on the set)
154 | export
155 | formGivenVals : (s : SortedSet _) -> List TTImp -> Vect s.size TTImp
156 | formGivenVals a b = fgvImpl (Vect.fromList $ Prelude.toList a) b
157 |   where
158 |     fgvImpl : Vect l _ -> List TTImp -> Vect l TTImp
159 |     fgvImpl []        _         = []
160 |     fgvImpl (_ :: xs) []        = `(_) :: fgvImpl xs []
161 |     fgvImpl (x :: xs) (y :: ys) = y    :: fgvImpl xs ys
162 |
163 | genGivens : List (TTImp, Fin x, Arg) -> (s : SortedSet (Fin x) ** Vect s.size TTImp)
164 | genGivens l = do
165 |   let (l1, l2, l3) = unzip3 l
166 |   let s = SortedSet.fromList l2
167 |   let gv = formGivenVals s l1
168 |   (s ** gv)
169 |
170 | -- Using the monadic trick makes the performance *much* better.
171 | specTaskToName : Monad m => TTImp -> m Name
172 | specTaskToName t = do
173 |   let (_, lamBody) = unLambda t
174 |   let (callee, _) = unAppAny lamBody
175 |   let vname =
176 |     case callee of
177 |          (IVar _ n) => show $ snd $ unNS n
178 |          x => show x
179 |   hash <- pure $ show $ hash t
180 |   pure $ fromString "\{vname}^\{hash}.\{vname}^\{hash}"
181 |
182 | nameUnambigAndVis : Elaboration m => Name -> m Bool
183 | nameUnambigAndVis n = do
184 |   try (do
185 |     _ : Unit <- check `(let x = ~(var n) in ())
186 |     pure True) (pure False)
187 |
188 | allConstructorsVisible : Elaboration m => TypeInfo -> m Bool
189 | allConstructorsVisible ti = do
190 |   all id <$> traverse (nameUnambigAndVis . name) ti.cons
191 |
192 | mkDPairOfUnknowns : Nat -> (Name -> TTImp) -> TTImp -> TTImp
193 | mkDPairOfUnknowns 0 _ t = t
194 | mkDPairOfUnknowns (S n) helper t = do
195 |   let nn = fromString $ "dph^\{show n}"
196 |   `(MkDPair ~(helper nn) ~(mkDPairOfUnknowns n helper t))
197 |
198 | dPairOfUnknowns : Nat -> TTImp
199 | dPairOfUnknowns 0 = `(?)
200 | dPairOfUnknowns (S n) = `(DPair ? $ \_ => ~(dPairOfUnknowns n))
201 |
202 | inSameNS : (nsSource: Name) -> Name -> Name
203 | inSameNS (NS ns _) n = NS ns n
204 | inSameNS _ n = n
205 |
206 | export
207 | specialiseIfNeeded :
208 |   Elaboration m =>
209 |   NamesInfoInTypes =>
210 |   ConsRecs =>
211 |   DerivationClosure m =>
212 |   (sig : GenSignature) ->
213 |   (fuel : TTImp) ->
214 |   Vect sig.givenParams.size TTImp ->
215 |   m $ Maybe TTImp
216 | specialiseIfNeeded sig fuel givenParamValues = do
217 |   logPoint DetailedDebug "deptycheck.derive.specialisation" [sig] "Checking specialisation need for \{show givenParamValues}..."
218 |   -- Check if there are any given type args, if not return Nothing
219 |   let True = any (\a => snd (unPi a.type) == `(Type)) $ index' sig.targetType.args <$> Prelude.toList sig.givenParams
220 |     | False =>
221 |       logValue DetailedDebug "deptycheck.derive.specialisation" [sig]
222 |         "Not found any given type args, specialisation not needed."
223 |         Nothing
224 |   -- Check if all of the generated type's constructors are visible, if not return Nothing
225 |   True <- allConstructorsVisible sig.targetType
226 |     | False =>
227 |       logValue DetailedDebug "deptycheck.derive.specialisation" [sig]
228 |         "\{sig.targetType.name} has invisible constructors, specialisation impossible."
229 |         Nothing
230 |   -- Assemble the `GenArg`s from `GenSignature` and given values
231 |   let givenIdxVals = Prelude.toList sig.givenParams `zipV` givenParamValues
232 |   let genArgs = mkArgs sig (withIndex sig.targetType.args) givenIdxVals
233 |   -- Check if at least one `GenArg` can be specialised upon (i.e. is a type argument and has a non-passthrough given value)
234 |   -- We need to terminate when all givens are passthrough, because otherwise we'll be stuck endlessly performing
235 |   -- identity specialisations of the same type
236 |   False <- all id <$> traverse (.isPassthrough) genArgs
237 |     | True =>
238 |       logValue DetailedDebug "deptycheck.derive.specialisation" [sig]
239 |         "Not found any type arguments that can be specialised upon, specialisation impossible."
240 |         Nothing
241 |   -- Generate specialisation rhs, arguments, and given values
242 |   (lambdaRet, fvArgs, givenSubst) <- processArgs sig genArgs
243 |   let preNorm = foldr lam lambdaRet fvArgs
244 |   logPoint DetailedDebug "deptycheck.derive.specialisation" [sig] "Task before normalisation: \{show preNorm}"
245 |   -- Normalise the specialisation lambda
246 |   (lambdaTy, lambdaBody) <- normaliseTask fvArgs lambdaRet
247 |   logPoint DetailedDebug "deptycheck.derive.specialisation" [sig] "NormaliseTask returned: lambdaTy = \{show lambdaTy};"
248 |   logPoint DetailedDebug "deptycheck.derive.specialisation" [sig] "                        lambdaBody = \{show lambdaBody};"
249 |   -- Generate specialised type name
250 |   specName <- specTaskToName lambdaBody
251 |   logPoint DetailedDebug "deptycheck.derive.specialisation" [sig] "Specialised type name: \{show specName}"
252 |   -- Check if `NamesInfoInTypes` contains specialised type
253 |   (specTy, specDecls) : (TypeInfo, List Decl) <- case lookupType specName of
254 |     -- If not, try looking it up via elaborator
255 |     Nothing => do
256 |       info <- try (Just <$> getInfo' specName) (pure Nothing)
257 |       case info of
258 |         Nothing => do
259 |         -- If not found at all, derive specialised type
260 |           logPoint DetailedDebug "deptycheck.derive.specialisation" [sig] "Specialised type not found, deriving..."
261 |           thisNS <- do
262 |             NS nsn _ <- inCurrentNS ""
263 |             | _ => fail "Internal error: inCurrentNS did not return NS"
264 |             pure nsn
265 |           Right (specTy, specDecls) <- runEitherT {m} {e=SpecialisationError} $
266 |               specialiseDataRaw {nsProvider = inNS thisNS} specName lambdaTy lambdaBody
267 |             | Left err => fail "INTERNAL ERROR: Specialisation \{show lambdaBody} failed with error \{show err}."
268 |           logPoint DetailedDebug "deptycheck.derive.specialisation" [sig] "Derived \{show specTy.name}"
269 |           -- Declare derived type
270 |           declare specDecls
271 |           specTy <- getInfo' specName
272 |           logValue Trace "deptycheck.derive.specialisation" [sig]
273 |             "Declared specialised type \{show specTy.name}: \{show lambdaRet}"
274 |             (specTy, [])
275 |         Just specTy =>
276 |           logValue DetailedDebug "deptycheck.derive.specialisation" [sig]
277 |             "Found \{show specTy.name}"
278 |             (specTy, [])
279 |     Just specTy =>
280 |       logValue DetailedDebug "deptycheck.derive.specialisation" [sig]
281 |         "Found \{show specTy.name}"
282 |         (specTy, [])
283 |   -- Assert that all of the specialised type's arguments are named for the specialised generator's `GenSignature` (this property should always be true)
284 |   let Yes stNamed = areAllTyArgsNamed specTy
285 |     | No _ => fail "INTERNAL ERROR: Specialised type \{show specTy.name} does not have fully named arguments and constructors."
286 |   -- Form new givens set and given value list
287 |   let (newGP ** newGVals) = genGivens $ mapMaybe (\(a,b) => map (,b) a) $ zip givenSubst $ withIndex specTy.args
288 |   -- Obtain the specialised generator call
289 |   (inv, cg_rhs) <- callGen (MkGenSignature specTy newGP) fuel newGVals
290 |   let inv : TTImp = case cg_rhs of
291 |         Nothing => inv
292 |         Just (n ** perm) => reorderGend False perm inv
293 |   -- Use derived cast to convert result back to polymorphic type
294 |   let generateds = sig.targetType.args.length `minus` sig.givenParams.size
295 |   let inv : TTImp =
296 |     if generateds == 0
297 |         then `(map (cast @{~(var $ inSameNS specTy.name "mToP")}) $ ~inv)
298 |         else
299 |           `(the (Gen MaybeEmpty ~(dPairOfUnknowns generateds)) $ map (\invv =>
300 |             case invv of
301 |               ~(mkDPairOfUnknowns generateds bindVar (bindVar "inv")) =>
302 |                   ~(mkDPairOfUnknowns generateds var `(cast @{~(var $ inSameNS specTy.name "mToP")} inv))) ~inv)
303 |   pure $ Just inv
304 |