fix: tcWholeProg updates metadata of ADT parameter selections (#1138)

primer/src/Primer/Action/Available.hs

@@ -128,6 +128,7 @@ import Primer.Zipper (
   findNodeWithParent,
   findType,
   focusOn,
+  focusOnKind,
   focusOnTy,
   locToEither,
   target,
@@ -423,13 +424,7 @@ forTypeDefParamKindNode paramName id l Editable tydefs defs tdName td =
       Just (KHole _) -> [NoInput MakeKType]
       Just _ -> [NoInput DeleteKind]
   where
-    findKind i k =
-      if getID k == i
-        then Just k
-        else case k of
-          KHole _ -> Nothing
-          KType _ -> Nothing
-          KFun _ k1 k2 -> findKind i k1 <|> findKind i k2
+    findKind i k = target <$> focusOnKind i k
 
 forTypeDefConsNode ::
   Level ->

primer/src/Primer/App.hs

@@ -86,12 +86,12 @@ import Optics (
   Field2 (_2),
   Field3 (_3),
   Fold,
+  afolding,
   elemOf,
   folded,
   getting,
   ifoldMap,
   mapped,
-  notElemOf,
   over,
   set,
   summing,
@@ -169,14 +169,14 @@ import Primer.Core (
   _chkedAt,
   _exprMeta,
   _exprMetaLens,
+  _kindMetaLens,
   _synthed,
   _type,
   _typeMetaLens,
  )
 import Primer.Core.DSL (S, create, emptyHole, kfun, khole, ktype, tEmptyHole)
 import Primer.Core.DSL qualified as DSL
 import Primer.Core.Transform (renameTyVar, renameVar, unfoldTApp)
-import Primer.Core.Type.Utils (kindIDs)
 import Primer.Core.Utils (freeVars, generateKindIDs, generateTypeIDs, regenerateExprIDs, regenerateTypeIDs, _freeTmVars, _freeTyVars, _freeVarsTy)
 import Primer.Def (
   ASTDef (..),
@@ -240,6 +240,7 @@ import Primer.Zipper (
   current,
   focusLoc,
   focusOn,
+  focusOnKind,
   focusOnTy,
   foldAbove,
   foldAboveTypeZ,
@@ -1026,16 +1027,9 @@ applyProgAction prog = \case
         >>= either (throwError . ActionError) pure
     pure (mods', Nothing)
     where
-      modifyKind f k
-        | notElemOf kindIDs id k = throwError' $ IDNotFound id
-        | otherwise = modifyKind' f k
-      modifyKind' f k =
-        if getID k == id
-          then f k
-          else case k of
-            KHole _ -> pure k
-            KType _ -> pure k
-            KFun m k1 k2 -> KFun m <$> modifyKind' f k1 <*> modifyKind' f k2
+      modifyKind f k = fromMaybe (throwError' $ IDNotFound id) $ do
+        k' <- focusOnKind id k
+        pure $ fromZipper . flip replace k' <$> f (target k')
       replaceHole a r = \case
         KHole{} -> r
         _ -> throwError' $ CustomFailure a "can only construct this kind in a hole"
@@ -1625,22 +1619,46 @@ tcWholeProg p = do
             { def = defName_
             , node = updatedNode
             }
-    Just (SelectionTypeDef s) ->
-      pure . Just . SelectionTypeDef $
-        s & over (#node % mapped % #_TypeDefConsNodeSelection) \conSel ->
-          conSel & over #field \case
+    Just (SelectionTypeDef s) -> do
+      let defName_ = s.def
+      -- If something goes wrong in finding the metadata, we just don't set a field selection.
+      -- This is similar to what we do when selection is in a term, above.
+      let updatedNode = case liftA2 (,) s.node (typeDefAST =<< Map.lookup s.def (allTypesMeta p)) of
             Nothing -> Nothing
-            Just fieldSel ->
-              flip (set #meta) fieldSel . (Right . Left . (^. _typeMetaLens)) <$> do
-                -- If something goes wrong in finding the metadata, we just don't set a field selection.
-                -- This is similar to what we do when selection is in a term, above.
-                td <- Map.lookup s.def $ allTypesMeta p
-                tda <- typeDefAST td
-                ty <- getTypeDefConFieldType tda conSel.con fieldSel.index
-                id <- case fieldSel.meta of
-                  Left _ -> Nothing -- Any selection in a typedef should have TypeMeta or KindMeta, not ExprMeta
-                  Right m -> pure $ getID m
-                target <$> focusOnTy id ty
+            Just (sn, tda) -> case sn of
+              TypeDefParamNodeSelection paramSel ->
+                Just $
+                  TypeDefParamNodeSelection $
+                    paramSel
+                      & #kindMeta %~ \case
+                        Just (Right (Right m)) -> do
+                          k <-
+                            tda
+                              ^? #astTypeDefParameters
+                              % afolding
+                                (find ((== paramSel.param) . fst))
+                              % _2
+                          Right . Right . view _kindMetaLens . target <$> focusOnKind (getID m) k
+                        _ -> Nothing
+              TypeDefConsNodeSelection conSel ->
+                Just $
+                  TypeDefConsNodeSelection $
+                    conSel & over #field \case
+                      Nothing -> Nothing
+                      Just fieldSel ->
+                        flip (set #meta) fieldSel . (Right . Left . (^. _typeMetaLens)) <$> do
+                          ty <- getTypeDefConFieldType tda conSel.con fieldSel.index
+                          id <- case fieldSel.meta of
+                            Left _ -> Nothing -- Any selection in a typedef should have TypeMeta or KindMeta, not ExprMeta
+                            Right m -> pure $ getID m
+                          target <$> focusOnTy id ty
+      pure $
+        Just $
+          SelectionTypeDef $
+            TypeDefSelection
+              { def = defName_
+              , node = updatedNode
+              }
   pure $ p'{progSelection = newSel}
 
 -- | Do a full check of a 'Prog', both the imports and the local modules

primer/src/Primer/Core.hs

@@ -84,6 +84,7 @@ import Primer.Core.Type (
   Type' (..),
   TypeMeta,
   _kindMeta,
+  _kindMetaLens,
   _typeMeta,
   _typeMetaLens,
  )

primer/src/Primer/Core/Type.hs

@@ -8,6 +8,7 @@ module Primer.Core.Type (
   _typeMetaLens,
   KindMeta,
   _kindMeta,
+  _kindMetaLens,
 ) where
 
 import Foreword
@@ -91,6 +92,12 @@ data Kind' a
 _kindMeta :: Traversal (Kind' a) (Kind' b) a b
 _kindMeta = param @0
 
+-- | A lens on to the metadata of a kind.
+-- Note that unlike '_kindMeta', this is shallow i.e. it does not recurse in to sub-expressions.
+-- And for this reason, it cannot be type-changing.
+_kindMetaLens :: Lens' (Kind' a) a
+_kindMetaLens = position @1
+
 instance HasID a => HasID (Type' a) where
   _id = position @1 % _id
 

primer/src/Primer/Zipper.hs

@@ -26,6 +26,7 @@ module Primer.Zipper (
   replace,
   focusOn,
   focusOnTy,
+  focusOnKind,
   top,
   up,
   down,
@@ -125,6 +126,7 @@ import Primer.Zipper.Type (
   down,
   farthest,
   focus,
+  focusOnKind,
   focusOnTy,
   foldAbove,
   foldBelow,

primer/src/Primer/Zipper/Type.hs

@@ -10,6 +10,7 @@ module Primer.Zipper.Type (
   target,
   _target,
   replace,
+  focusOnKind,
   focusOnTy,
   top,
   up,
@@ -58,11 +59,14 @@ import Primer.Core.Meta (
   unLocalName,
  )
 import Primer.Core.Type (
+  Kind',
   Type' (TForall, TLet),
   TypeMeta,
  )
 import Primer.Name (Name)
 
+type KindZip' c = Zipper (Kind' c) (Kind' c)
+
 type TypeZip' b = Zipper (Type' b) (Type' b)
 
 -- | An ordinary zipper for 'Type's
@@ -109,6 +113,27 @@ focus = zipper
 replace :: (IsZipper za a) => a -> za -> za
 replace = over asZipper . replaceHole
 
+-- | Focus on the node with the given 'ID', if it exists in the kind
+focusOnKind ::
+  (Data c, HasID c) =>
+  ID ->
+  Kind' c ->
+  Maybe (KindZip' c)
+focusOnKind i = focusOnKind' i . focus
+
+-- | Focus on the node with the given 'ID', if it exists in the focussed kind
+focusOnKind' ::
+  (Data c, HasID c) =>
+  ID ->
+  KindZip' c ->
+  Maybe (KindZip' c)
+focusOnKind' i = fmap snd . search matchesID
+  where
+    matchesID z
+      -- If the current target has the correct ID, return that
+      | getID (target z) == i = Just z
+      | otherwise = Nothing
+
 -- | Focus on the node with the given 'ID', if it exists in the type
 focusOnTy ::
   (Data b, HasID b) =>