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 Deriving.SpecialiseData.TaskFormation
 15 |
 16 | import public Language.Reflection.Unify
 17 |
 18 | import public Data.Hashable
 19 | import public Data.Hashable.Base
 20 |
 21 | %default total
 22 |
 23 | ||| Determine if the argument should be specialised or passed through
 24 | (.isPassthrough) : Elaboration m => GenArg -> m Bool
 25 | (.isPassthrough) (MkGenArg a Nothing) = pure True
 26 | (.isPassthrough) (MkGenArg a $ Just g) = do
 27 |   let True = snd (unPi a.type) == `(Type)
 28 |     | _ => pure True
 29 |   case g of
 30 |     IVar _ n => do
 31 |       nInfo <- getInfo n
 32 |       case nInfo of
 33 |         [] => pure True
 34 |         _ => pure False
 35 |     _ => pure False
 36 |
 37 | ||| Assemble a list of arguments and their given values from `callGen` inputs
 38 | |||
 39 | ||| The indices inside both given lists must be in ascending order
 40 | mkArgs :
 41 |   NamesInfoInTypes =>
 42 |   (sig : GenSignature) ->
 43 |   List (Fin sig.targetType.args.length, Arg) ->
 44 |   List (Fin sig.targetType.args.length, TTImp) ->
 45 |   List GenArg
 46 | mkArgs sig [] _ = []
 47 | mkArgs sig ((_, x) :: xs) [] = MkGenArg x Nothing :: mkArgs sig xs []
 48 | mkArgs sig ((i1, x) :: xs) g@((i2, y) :: ys) =
 49 |   if i1 == i2
 50 |     then MkGenArg x (Just y) :: mkArgs sig xs ys
 51 |     else MkGenArg x Nothing  :: mkArgs sig xs g
 52 |
 53 | ||| Given a set of given argument indices, convert a list of their values into a vector that can be fed to `callGen`
 54 | |||
 55 | ||| The values should be listed for indices in ascending order
 56 | ||| (i.e. how these indices would be sorted if we called `toList` on the set)
 57 | export
 58 | formGivenVals : (s : SortedSet _) -> List TTImp -> Vect s.size TTImp
 59 | formGivenVals a b = fgvImpl (Vect.fromList $ Prelude.toList a) b
 60 |   where
 61 |     fgvImpl : Vect l _ -> List TTImp -> Vect l TTImp
 62 |     fgvImpl []        _         = []
 63 |     fgvImpl (_ :: xs) []        = `(_) :: fgvImpl xs []
 64 |     fgvImpl (x :: xs) (y :: ys) = y    :: fgvImpl xs ys
 65 |
 66 | genGivens : List (TTImp, Fin x, Arg) -> (s : SortedSet (Fin x) ** Vect s.size TTImp)
 67 | genGivens l = do
 68 |   let (l1, l2, l3) = unzip3 l
 69 |   let s = SortedSet.fromList l2
 70 |   let gv = formGivenVals s l1
 71 |   (s ** gv)
 72 |
 73 | -- Using the monadic trick makes the performance *much* better.
 74 | specTaskToName : Monad m => TTImp -> m Name
 75 | specTaskToName t = do
 76 |   let (_, lamBody) = unLambda t
 77 |   let (callee, _) = unAppAny lamBody
 78 |   let vname =
 79 |     case callee of
 80 |          (IVar _ n) => show $ snd $ unNS n
 81 |          x => show x
 82 |   hash <- pure $ show $ hash t
 83 |   pure $ fromString "\{vname}^\{hash}.\{vname}^\{hash}"
 84 |
 85 | nameUnambigAndVis : Elaboration m => Name -> m Bool
 86 | nameUnambigAndVis n = do
 87 |   try (do
 88 |     _ : Unit <- check `(let x = ~(var n) in ())
 89 |     pure True) (pure False)
 90 |
 91 | allConstructorsVisible : Elaboration m => TypeInfo -> m Bool
 92 | allConstructorsVisible ti = do
 93 |   all id <$> traverse (nameUnambigAndVis . name) ti.cons
 94 |
 95 | mkDPairOfUnknowns : Nat -> (Name -> TTImp) -> TTImp -> TTImp
 96 | mkDPairOfUnknowns 0 _ t = t
 97 | mkDPairOfUnknowns (S n) helper t = do
 98 |   let nn = fromString $ "dph^\{show n}"
 99 |   `(MkDPair ~(helper nn) ~(mkDPairOfUnknowns n helper t))
100 |
101 | dPairOfUnknowns : Nat -> TTImp
102 | dPairOfUnknowns 0 = `(?)
103 | dPairOfUnknowns (S n) = `(DPair ? $ \_ => ~(dPairOfUnknowns n))
104 |
105 | inSameNS : (nsSource: Name) -> Name -> Name
106 | inSameNS (NS ns _) n = NS ns n
107 | inSameNS _ n = n
108 |
109 | export
110 | specialiseIfNeeded :
111 |   Elaboration m =>
112 |   NamesInfoInTypes =>
113 |   ConsRecs =>
114 |   DerivationClosure m =>
115 |   (sig : GenSignature) ->
116 |   (fuel : TTImp) ->
117 |   Vect sig.givenParams.size TTImp ->
118 |   m $ Maybe TTImp
119 | specialiseIfNeeded sig fuel givenParamValues = do
120 |   logPoint DetailedDebug "deptycheck.derive.specialisation" [sig] "Checking specialisation need for \{show givenParamValues}..."
121 |   -- Check if there are any given type args, if not return Nothing
122 |   let True = any (\a => snd (unPi a.type) == `(Type)) $ index' sig.targetType.args <$> Prelude.toList sig.givenParams
123 |     | False =>
124 |       logValue DetailedDebug "deptycheck.derive.specialisation" [sig]
125 |         "Not found any given type args, specialisation not needed."
126 |         Nothing
127 |   -- Check if all of the generated type's constructors are visible, if not return Nothing
128 |   True <- allConstructorsVisible sig.targetType
129 |     | False =>
130 |       logValue DetailedDebug "deptycheck.derive.specialisation" [sig]
131 |         "\{sig.targetType.name} has invisible constructors, specialisation impossible."
132 |         Nothing
133 |   -- Assemble the `GenArg`s from `GenSignature` and given values
134 |   let givenIdxVals = Prelude.toList sig.givenParams `zipV` givenParamValues
135 |   let genArgs = mkArgs sig (withIndex sig.targetType.args) givenIdxVals
136 |   -- Check if at least one `GenArg` can be specialised upon (i.e. is a type argument and has a non-passthrough given value)
137 |   -- We need to terminate when all givens are passthrough, because otherwise we'll be stuck endlessly performing
138 |   -- identity specialisations of the same type
139 |   False <- all id <$> traverse (.isPassthrough) genArgs
140 |     | True =>
141 |       logValue DetailedDebug "deptycheck.derive.specialisation" [sig]
142 |         "Not found any type arguments that can be specialised upon, specialisation impossible."
143 |         Nothing
144 |   -- Generate specialisation rhs, arguments, and given values
145 |   (lambdaRet, fvArgs, givenSubst) <- argsToSpecTask specDecideArg sig.targetType.name genArgs
146 |   let preNorm = foldr lam lambdaRet fvArgs
147 |   logPoint DetailedDebug "deptycheck.derive.specialisation" [sig] "Task before normalisation: \{show preNorm}"
148 |   -- Normalise the specialisation lambda
149 |   (lambdaTy, lambdaBody) <- normaliseTask fvArgs lambdaRet
150 |   logPoint DetailedDebug "deptycheck.derive.specialisation" [sig] "NormaliseTask returned: lambdaTy = \{show lambdaTy};"
151 |   logPoint DetailedDebug "deptycheck.derive.specialisation" [sig] "                        lambdaBody = \{show lambdaBody};"
152 |   -- Generate specialised type name
153 |   specName <- specTaskToName lambdaBody
154 |   logPoint DetailedDebug "deptycheck.derive.specialisation" [sig] "Specialised type name: \{show specName}"
155 |   -- Check if `NamesInfoInTypes` contains specialised type
156 |   (specTy, specDecls) : (TypeInfo, List Decl) <- case lookupType specName of
157 |     -- If not, try looking it up via elaborator
158 |     Nothing => do
159 |       info <- try (Just <$> getInfo' specName) (pure Nothing)
160 |       case info of
161 |         Nothing => do
162 |         -- If not found at all, derive specialised type
163 |           logPoint DetailedDebug "deptycheck.derive.specialisation" [sig] "Specialised type not found, deriving..."
164 |           thisNS <- do
165 |             NS nsn _ <- inCurrentNS ""
166 |             | _ => fail "Internal error: inCurrentNS did not return NS"
167 |             pure nsn
168 |           Right (specTy, specDecls) <- runEitherT {m} {e=SpecialisationError} $
169 |               specialiseDataRaw {nsProvider = inNS thisNS} specName lambdaTy lambdaBody
170 |             | Left err => fail "INTERNAL ERROR: Specialisation \{show lambdaBody} failed with error \{show err}."
171 |           logPoint DetailedDebug "deptycheck.derive.specialisation" [sig] "Derived \{show specTy.name}"
172 |           -- Declare derived type
173 |           declare specDecls
174 |           specTy <- getInfo' specName
175 |           logValue Trace "deptycheck.derive.specialisation" [sig]
176 |             "Declared specialised type \{show specTy.name}: \{show lambdaRet}"
177 |             (specTy, [])
178 |         Just specTy =>
179 |           logValue DetailedDebug "deptycheck.derive.specialisation" [sig]
180 |             "Found \{show specTy.name}"
181 |             (specTy, [])
182 |     Just specTy =>
183 |       logValue DetailedDebug "deptycheck.derive.specialisation" [sig]
184 |         "Found \{show specTy.name}"
185 |         (specTy, [])
186 |   -- Assert that all of the specialised type's arguments are named for the specialised generator's `GenSignature` (this property should always be true)
187 |   let Yes stNamed = areAllTyArgsNamed specTy
188 |     | No _ => fail "INTERNAL ERROR: Specialised type \{show specTy.name} does not have fully named arguments and constructors."
189 |   -- Form new givens set and given value list
190 |   let (newGP ** newGVals) = genGivens $ mapMaybe (\(a,b) => map (,b) a) $ zip givenSubst $ withIndex specTy.args
191 |   -- Obtain the specialised generator call
192 |   (inv, cg_rhs) <- callGen (MkGenSignature specTy newGP) fuel newGVals
193 |   let inv : TTImp = case cg_rhs of
194 |         Nothing => inv
195 |         Just (n ** perm) => reorderGend False perm inv
196 |   -- Use derived cast to convert result back to polymorphic type
197 |   let generateds = sig.targetType.args.length `minus` sig.givenParams.size
198 |   let inv : TTImp =
199 |     if generateds == 0
200 |         then `(map (cast @{~(var $ inSameNS specTy.name "mToP")}) $ ~inv)
201 |         else
202 |           `(the (Gen MaybeEmpty ~(dPairOfUnknowns generateds)) $ map (\invv =>
203 |             case invv of
204 |               ~(mkDPairOfUnknowns generateds bindVar (bindVar "inv")) =>
205 |                   ~(mkDPairOfUnknowns generateds var `(cast @{~(var $ inSameNS specTy.name "mToP")} inv))) ~inv)
206 |   pure $ Just inv
207 |