|
| 1 | +module Deriving.DepTyCheck.Util.Specialisation |
| 2 | + |
| 3 | +import public Control.Monad.Either |
| 4 | + |
| 5 | +import public Data.DPair |
| 6 | +import public Data.List.Ex |
| 7 | +import public Data.List.Map |
| 8 | +import public Data.SortedMap |
| 9 | +import public Data.SortedMap.Extra |
| 10 | +import public Data.SortedSet |
| 11 | + |
| 12 | +import public Deriving.DepTyCheck.Gen.ForOneType.Interface |
| 13 | + |
| 14 | +import public Deriving.SpecialiseData |
| 15 | +import public Language.Reflection.Unify |
| 16 | + |
| 17 | +%default total |
| 18 | + |
| 19 | +isSimpleN : Name -> Bool |
| 20 | +isSimpleN (UN (Basic str)) = (fst $ span isLower str) /= "" |
| 21 | +isSimpleN _ = False |
| 22 | + |
| 23 | +||| Given an argument, a possible given value for it, and a set of free variable names in scope, |
| 24 | +||| return argument value in the specialisation lambda body, |
| 25 | +||| as well as a list of newly introduced free variables and their given values |
| 26 | +processArg : Arg -> Maybe TTImp -> SortedSet Name -> (TTImp, List (Arg, Maybe TTImp)) |
| 27 | +processArg arg Nothing fvNames = (IVar EmptyFC $ fromMaybe "_" arg.name, [(arg, Nothing)]) |
| 28 | +processArg arg (Just given) fvNames = |
| 29 | + case snd $ unPi arg.type of |
| 30 | + `(Type) => do |
| 31 | + let simpleNames = filter isSimpleN $ allVarNames given |
| 32 | + let newNames = filter (not . contains' fvNames) simpleNames |
| 33 | + let newArgs : List (Arg, Maybe TTImp) = newNames <&> \x => (MkArg MW ExplicitArg (Just x) `(_), Just $ IVar EmptyFC x) |
| 34 | + (given, newArgs) |
| 35 | + _ => (IVar EmptyFC $ fromMaybe "_" arg.name, [(arg, Just given)]) |
| 36 | + |
| 37 | +processArgs' : |
| 38 | + List (Fin x, Arg) -> |
| 39 | + List (Fin x, TTImp) -> |
| 40 | + SnocList (Arg, Maybe TTImp) -> |
| 41 | + SortedSet Name -> |
| 42 | + (List AnyApp, SnocList (Arg, Maybe TTImp), SortedSet Name) |
| 43 | +processArgs' [] _ fvArgs fvNames = ([], fvArgs, fvNames) |
| 44 | +processArgs' ((argIdx, arg) :: xs) givens fvArgs fvNames = do |
| 45 | + let (givens, argVal, newFVs) : (List _, TTImp, List _) = case givens of |
| 46 | + [] => ([], processArg arg Nothing fvNames) |
| 47 | + (givenIdx, givenVal) :: ys => |
| 48 | + if givenIdx == argIdx |
| 49 | + then (ys, processArg arg (Just givenVal) fvNames) |
| 50 | + else (givens, processArg arg Nothing fvNames) |
| 51 | + let newFVNames = SortedSet.fromList $ fromMaybe "_" . name . fst <$> newFVs |
| 52 | + let fvArgs = fvArgs <>< newFVs |
| 53 | + let fvNames = union fvNames newFVNames |
| 54 | + mapFst (appArg arg argVal ::) $ processArgs' xs givens fvArgs fvNames |
| 55 | + |
| 56 | +processArgs : |
| 57 | + (sig : GenSignature) -> |
| 58 | + List (Fin sig.targetType.args.length, Arg) -> |
| 59 | + List (Fin sig.targetType.args.length, TTImp) -> |
| 60 | + (TTImp, List Arg, List $ Maybe TTImp) |
| 61 | +processArgs sig args givens = do |
| 62 | + let (argVals, fvArgs, _) = processArgs' args givens [<] empty |
| 63 | + (reAppAny (IVar EmptyFC sig.targetType.name) argVals, unzip $ toList fvArgs) |
| 64 | + |
| 65 | +remap : List (Maybe TTImp, Fin x, Arg) -> List (TTImp, Fin x, Arg) |
| 66 | +remap [] = [] |
| 67 | +remap ((Nothing, _, _) :: xs) = remap xs |
| 68 | +remap ((Just x, y, z) :: xs) = (x,y,z) :: remap xs |
| 69 | + |
| 70 | +formGivenVals : Vect l _ -> List TTImp -> Vect l TTImp |
| 71 | +formGivenVals [] _ = [] |
| 72 | +formGivenVals (_ :: xs) [] = (`(_) :: formGivenVals xs []) |
| 73 | +formGivenVals (x :: xs) (y :: ys) = y :: formGivenVals xs ys |
| 74 | + |
| 75 | +genSS : List (TTImp, Fin x, Arg) -> (s : SortedSet (Fin x) ** Vect s.size TTImp) |
| 76 | +genSS l = do |
| 77 | + let (l1, l2) = unzip l |
| 78 | + let (l2, l3) = unzip l2 |
| 79 | + let s = SortedSet.fromList l2 |
| 80 | + let gv = formGivenVals (Vect.fromList $ Prelude.toList s) l1 |
| 81 | + (s ** gv) |
| 82 | + |
| 83 | +%tcinline |
| 84 | +specialiseIfNeeded : |
| 85 | + Monad m => |
| 86 | + Elaboration m => |
| 87 | + NamesInfoInTypes => |
| 88 | + DerivationClosure m => |
| 89 | + (sig : GenSignature) -> |
| 90 | + (specTaskToName : TTImp -> Name) -> |
| 91 | + (fuel : TTImp) -> |
| 92 | + Vect sig.givenParams.size TTImp -> |
| 93 | + m $ Maybe (List Decl, TypeInfo, TTImp) |
| 94 | +specialiseIfNeeded sig specTaskToName fuel givenParamValues = do |
| 95 | + -- Check if there are any given type args, if not return Nothing |
| 96 | + let True = any (\a => snd (unPi a.type) == `(Type)) $ index' sig.targetType.args <$> Prelude.toList sig.givenParams |
| 97 | + | False => pure Nothing |
| 98 | + let givenIdxVals = (Prelude.toList sig.givenParams) `zipV` givenParamValues |
| 99 | + let (lambdaRet, fvArgs, givenSubst) = processArgs sig (withIndex sig.targetType.args) givenIdxVals |
| 100 | + (lambdaTy, lambdaBody) <- normaliseTask fvArgs lambdaRet |
| 101 | + let specName = specTaskToName lambdaBody |
| 102 | + (specTy, specDecls) : (TypeInfo, List Decl) <- case lookupType specName of |
| 103 | + Nothing => do |
| 104 | + Right (specTy, specDecls) <- runEitherT {m} {e=SpecialisationError} $ specialiseDataRaw specName lambdaTy lambdaBody |
| 105 | + | Left err => fail "INTERNAL ERROR: Specialisation \{show lambdaBody} failed with error \{show err}." |
| 106 | + pure (specTy, specDecls) |
| 107 | + Just specTy => pure (specTy, []) |
| 108 | + let Yes stNamed = areAllTyArgsNamed specTy |
| 109 | + | No _ => fail "INTERNAL ERROR: Specialised type \{show specTy.name} does not have fully named arguments and constructors." |
| 110 | + let (newGP ** newGVals) = genSS $ remap $ zip givenSubst $ withIndex specTy.args |
| 111 | + (inv, cg_rhs) <- callGen (MkGenSignature specTy newGP) fuel newGVals |
| 112 | + let inv = case cg_rhs of |
| 113 | + Nothing => inv |
| 114 | + Just (n ** perm) => reorderGend False perm inv |
| 115 | + pure $ Just (specDecls, specTy, inv) |
0 commit comments