refactor: ParamKindAction uses existing machinery

Signed-off-by: Ben Price <[email protected]>

primer-api/src/Primer/API.hs

@@ -106,9 +106,9 @@ import Data.Tuple.Extra (curry3)
 import Optics (ifoldr, over, preview, to, traverseOf, view, (%), (^.), _Just)
 import Primer.API.NodeFlavor qualified as Flavor
 import Primer.API.RecordPair (RecordPair (RecordPair))
-import Primer.Action (ActionError, ProgAction, toProgActionInput, toProgActionNoInput)
+import Primer.Action (ActionError (ParamNotFound), ProgAction, toProgActionInput, toProgActionNoInput)
 import Primer.Action.Available qualified as Available
-import Primer.Action.ProgError (ProgError (NodeIDNotFound, ParamNotFound, TypeDefConFieldNotFound))
+import Primer.Action.ProgError (ProgError (ActionError, NodeIDNotFound, TypeDefConFieldNotFound))
 import Primer.App (
   App,
   DefSelection (..),
@@ -1355,7 +1355,7 @@ getSelectionTypeOrKind = curry $ logAPI (noError GetTypeOrKind) $ \(sid, sel0) -
         -- type def itself selected - return its kind
         Nothing -> pure $ Kind $ viewTreeKind' $ mkIdsK $ typeDefKind $ forgetTypeDefMetadata $ TypeDef.TypeDefAST def
         Just (TypeDefParamNodeSelection (TypeDefParamSelection p s)) -> do
-          k <- maybe (throw' $ ParamNotFound p) (pure . snd) $ find ((== p) . fst) (astTypeDefParameters def)
+          k <- maybe (throw' $ ActionError $ ParamNotFound p) (pure . snd) $ find ((== p) . fst) (astTypeDefParameters def)
           case s of
             Nothing ->
               -- param name node selected - return its kind

primer/src/Primer/Action.hs

@@ -20,6 +20,7 @@ module Primer.Action (
   uniquifyDefName,
   toProgActionInput,
   toProgActionNoInput,
+  applyActionsToParam,
   applyActionsToField,
   insertSubseqBy,
 ) where
@@ -32,14 +33,30 @@ import Data.Bifunctor.Swap qualified as Swap
 import Data.Bitraversable (bisequence)
 import Data.Functor.Compose (Compose (..))
 import Data.Generics.Product (typed)
+import Data.Generics.Uniplate.Zipper (fromZipper)
 import Data.List (delete, findIndex, insertBy)
 import Data.List.NonEmpty qualified as NE
 import Data.Map (insert)
 import Data.Map.Strict qualified as Map
 import Data.Set qualified as Set
 import Data.Text qualified as T
 import Data.Tuple.Extra ((&&&))
-import Optics (over, set, traverseOf, (%), (?~), (^.), (^?), _Just)
+import Optics (
+  findOf,
+  folded,
+  indices,
+  isnd,
+  over,
+  set,
+  traverseOf,
+  traversed,
+  (%),
+  (%&),
+  (?~),
+  (^.),
+  (^?),
+  _Just,
+ )
 import Primer.Action.Actions (Action (..), BranchMove (Fallback, Pattern), Movement (..), QualifiedText)
 import Primer.Action.Available qualified as Available
 import Primer.Action.Errors (ActionError (..))
@@ -64,6 +81,7 @@ import Primer.Core (
   HasID,
   HasMetadata (_metadata),
   ID,
+  Kind,
   Kind' (KHole),
   KindMeta,
   LVarName,
@@ -146,6 +164,7 @@ import Primer.Typecheck (
   checkEverything,
   exprTtoExpr,
   getTypeDefInfo,
+  initialCxt,
   lookupConstructor,
   lookupVar,
   maybeTypeOf,
@@ -168,6 +187,7 @@ import Primer.Zipper (
   findNodeWithParent,
   findTypeOrKind,
   focus,
+  focusKind,
   focusLoc,
   focusOn,
   focusOnlyKind,
@@ -186,6 +206,7 @@ import Primer.Zipper (
   updateCaseBind,
   _target,
  )
+import Primer.Zipper.Type (KindZip, focusOnlyKindT)
 import Primer.ZipperCxt (localVariablesInScopeExpr)
 
 -- | Given a definition name and a program, return a unique variant of
@@ -268,6 +289,40 @@ applyActionsToTypeSig smartHoles imports (mod, mods) (defName, def) actions =
             -- In this case we just refocus on the top of the type.
             z -> maybe unwrapError (pure . Left . focusOnlyType) (focusType (unfocusLoc z))
 
+applyActionsToParam ::
+  (MonadFresh ID m, MonadFresh NameCounter m) =>
+  SmartHoles ->
+  (TyVarName, ASTTypeDef TypeMeta KindMeta) ->
+  [Action] ->
+  m (Either ActionError (ASTTypeDef TypeMeta KindMeta, KindZip))
+applyActionsToParam sh (paramName, def) actions = runExceptT $ do
+  zk <- case findOf (#astTypeDefParameters % folded) ((== paramName) . fst) def of
+    Nothing -> throwError $ ParamNotFound paramName
+    Just (_, k) ->
+      -- no action in kinds should care about the context
+      flip runReaderT (initialCxt sh) $
+        withWrappedKind k $ \zk' ->
+          foldlM (flip applyActionAndSynth) (InKind zk') actions
+  let def' =
+        set
+          (#astTypeDefParameters % traversed % isnd %& indices (== paramName))
+          (fromZipper zk)
+          def
+  pure (def', zk)
+  where
+    withWrappedKind :: (MonadError ActionError m, MonadFresh ID m) => Kind -> (KindZ -> m Loc) -> m KindZip
+    withWrappedKind k f = do
+      wrappedKind <- ann emptyHole (tforall "a" (pure k) tEmptyHole)
+      let unwrapError = throwError $ InternalFailure "applyActionsToParam: failed to unwrap kind"
+          wrapError = throwError $ InternalFailure "applyActionsToParam: failed to wrap kind"
+          focusedKind = focusKind <=< focusType $ focus wrappedKind
+      case focusedKind of
+        Nothing -> wrapError
+        Just wrappedK ->
+          f wrappedK >>= \case
+            InKind zk -> pure $ focusOnlyKindT $ focusOnlyKind zk
+            z -> maybe unwrapError pure (fmap (focusOnlyKindT . focusOnlyKind) . focusKind <=< focusType $ unfocusLoc z)
+
 applyActionsToField ::
   (MonadFresh ID m, MonadFresh NameCounter m) =>
   SmartHoles ->

primer/src/Primer/Action/Errors.hs

@@ -13,7 +13,7 @@ import Data.Aeson (FromJSON (..), ToJSON (..))
 import Primer.Action.Actions (Action)
 import Primer.Action.Available qualified as Available
 import Primer.Action.Movement (Movement)
-import Primer.Core (Expr, GVarName, ID, LVarName, ModuleName, Pattern, TyConName, Type', ValConName)
+import Primer.Core (Expr, GVarName, ID, LVarName, ModuleName, Pattern, TyConName, TyVarName, Type', ValConName)
 import Primer.JSON (CustomJSON (..), PrimerJSON)
 import Primer.Typecheck.TypeError (TypeError)
 import Primer.Zipper (SomeNode)
@@ -79,6 +79,7 @@ data ActionError
   | NeedTypeDefParamKindSelection
   | NoNodeSelection
   | ValConNotFound TyConName ValConName
+  | ParamNotFound TyVarName
   | FieldIndexOutOfBounds ValConName Int
   deriving stock (Eq, Show, Read, Generic)
   deriving (FromJSON, ToJSON) via PrimerJSON ActionError

primer/src/Primer/Action/ProgError.hs

@@ -33,7 +33,6 @@ data ProgError
   | -- | We expected to see more arguments to a constructor than actually existed
     -- (this should never happen in a well-typed program)
     ConNotSaturated ValConName
-  | ParamNotFound TyVarName
   | NodeIDNotFound ID
   | TypeDefConFieldNotFound TyConName ValConName Int
   | ValConParamClash Name

primer/src/Primer/App.hs

@@ -117,6 +117,7 @@ import Primer.Action (
   applyAction',
   applyActionsToBody,
   applyActionsToField,
+  applyActionsToParam,
   applyActionsToTypeSig,
   insertSubseqBy,
  )
@@ -177,7 +178,7 @@ import Primer.Core (
   _type,
   _typeMetaLens,
  )
-import Primer.Core.DSL (S, create, emptyHole, kfun, khole, ktype, tEmptyHole)
+import Primer.Core.DSL (S, create, emptyHole, tEmptyHole)
 import Primer.Core.DSL qualified as DSL
 import Primer.Core.Transform (renameTyVar, renameVar, unfoldTApp)
 import Primer.Core.Utils (
@@ -779,7 +780,7 @@ applyProgAction prog = \case
         def
           & traverseOf
             #astTypeDefParameters
-            ( maybe (throwError $ ParamNotFound old) pure
+            ( maybe (throwError $ ActionError $ ParamNotFound old) pure
                 . findAndAdjust ((== old) . fst) (_1 .~ new)
             )
       updateConstructors =
@@ -947,7 +948,7 @@ applyProgAction prog = \case
               ( \ps -> do
                   unless
                     (paramName `elem` map fst ps)
-                    (throwError $ ParamNotFound paramName)
+                    (throwError $ ActionError $ ParamNotFound paramName)
                   pure $ filter ((/= paramName) . fst) ps
               )
               td
@@ -1021,42 +1022,21 @@ applyProgAction prog = \case
                   }
           )
   ParamKindAction tyName paramName id actions -> editModuleOfCrossType (Just tyName) prog $ \(mod, mods) defName def -> do
-    def' <-
-      def
-        & traverseOf
-          #astTypeDefParameters
-          ( maybe (throwError $ ParamNotFound paramName) pure
-              <=< findAndAdjustA
-                ((== paramName) . fst)
-                ( traverseOf _2 $
-                    flip
-                      ( foldlM $ flip \case
-                          ConstructKType -> modifyKind $ replaceHole ConstructKType ktype
-                          ConstructKFun -> modifyKind \k -> ktype `kfun` pure k
-                          Delete -> modifyKind $ const khole
-                          a -> const $ throwError $ ActionError $ CustomFailure a "unexpected non-kind action"
-                      )
-                      actions
-                )
-          )
-    let mod' = mod & over #moduleTypes (Map.insert defName $ TypeDefAST def')
-        imports = progImports prog
-        smartHoles = progSmartHoles prog
-    mods' <-
-      runExceptT
-        ( runReaderT
-            (checkEverything smartHoles (CheckEverything{trusted = imports, toCheck = mod' : mods}))
-            (buildTypingContextFromModules (mod : mods <> imports) smartHoles)
-        )
-        >>= either (throwError . ActionError) pure
-    pure (mods', Nothing)
-    where
-      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"
+    let smartHoles = progSmartHoles prog
+    res <- applyActionsToParam smartHoles (paramName, def) $ SetCursor id : actions
+    case res of
+      Left err -> throwError $ ActionError err
+      Right (def', _) -> do
+        let mod' = mod & over #moduleTypes (Map.insert defName $ TypeDefAST def')
+            imports = progImports prog
+        mods' <-
+          runExceptT
+            ( runReaderT
+                (checkEverything smartHoles (CheckEverything{trusted = imports, toCheck = mod' : mods}))
+                (buildTypingContextFromModules (mod : mods <> imports) smartHoles)
+            )
+            >>= either (throwError . ActionError) pure
+        pure (mods', Nothing)
   SetSmartHoles smartHoles ->
     pure $ prog & #progSmartHoles .~ smartHoles
   CopyPasteSig fromIds setup -> case mdefName of

primer/src/Primer/Zipper.hs

@@ -19,6 +19,7 @@ module Primer.Zipper (
   BindLoc,
   BindLoc' (..),
   focusType,
+  focusKind,
   focusLoc,
   unfocusType,
   unfocusKind,
@@ -114,7 +115,7 @@ import Primer.Core (
   LVarName,
   LocalName (unLocalName),
   Type,
-  Type' (),
+  Type' (TForall),
   TypeMeta,
   bindName,
   getID,
@@ -285,6 +286,22 @@ focusType z = case target z of
   where
     l = _target % typesInExpr
 
+-- | Switch from an 'Type' zipper to a 'Kind' zipper, focusing on the kind in
+-- the current target. This expects that the target is an @TForall@ node
+-- (as this is the only one that contain a @Kind@).
+focusKind :: (Data b, Data c) => TypeZ' a b c -> Maybe (KindZ' a b c)
+focusKind (ZipNest z f) = case target z of
+  TForall m n k t ->
+    pure $
+      ZipNest
+        ( ZipNest
+            (focus k)
+            $ \k' -> replace (TForall m n k' t) z
+        )
+        f
+  -- pure $ ZipNest (zipper t) $ \t' -> z & l .~ t'
+  _ -> Nothing
+
 -- | If the currently focused expression is a case expression, search the bindings of its branches
 -- to find one matching the given ID, and return the 'Loc' for that binding.
 -- If no match is found, return @Nothing@.

primer/src/Primer/Zipper/Type.hs

@@ -9,6 +9,7 @@ module Primer.Zipper.Type (
   KindTZ',
   KindTZ,
   unfocusKindT,
+  focusOnlyKindT,
   focusOnTy,
   focusOnTy',
   farthest,
@@ -55,6 +56,7 @@ import Primer.Zipper.Nested (
   ZipNest (ZipNest),
   down,
   focus,
+  innerZipNest,
   left,
   replace,
   right,
@@ -82,6 +84,10 @@ type KindTZ = KindTZ' TypeMeta KindMeta
 unfocusKindT :: Data c => KindTZ' b c -> TypeZip' b c
 unfocusKindT = unfocusNest
 
+-- | Forget the surrounding type context
+focusOnlyKindT :: KindTZ' b c -> KindZip' c
+focusOnlyKindT = innerZipNest
+
 -- | Focus on the node with the given 'ID', if it exists in the kind
 focusOnKind ::
   (Data c, HasID c) =>

primer/test/Tests/Action/Prog.hs

@@ -1198,17 +1198,17 @@ unit_ParamKindAction_2 =
     ( defaultProgEditableTypeDefs (pure [])
     )
     [ ParamKindAction tT pB 30 [ConstructKFun]
-    , ParamKindAction tT pB 32 [ConstructKType]
+    , ParamKindAction tT pB 36 [ConstructKType]
     ]
-    $ expectError (@?= ActionError (CustomFailure ConstructKType "can only construct this kind in a hole"))
+    $ expectError (@?= ActionError (CustomFailure ConstructKType "can only construct the kind 'Type' in hole"))
 
 unit_ParamKindAction_2b :: Assertion
 unit_ParamKindAction_2b =
   progActionTest
     ( defaultProgEditableTypeDefs (pure [])
     )
     [ ParamKindAction tT pB 30 [ConstructKFun]
-    , ParamKindAction tT pB 32 [Delete]
+    , ParamKindAction tT pB 36 [Delete]
     ]
     $ expectSuccess
     $ \_ prog' -> do