feat: Add actions for modifying constructor field types

Signed-off-by: George Thomas <[email protected]>

primer/src/Foreword.hs

@@ -32,6 +32,8 @@ module Foreword (
   curry4,
   unsafeMaximum,
   spanMaybe,
+  adjustAtA',
+  findAndAdjustA',
 ) where
 
 -- In general, we should defer to "Protolude"'s exports and avoid name
@@ -130,6 +132,11 @@ adjustAtA n f xs = case splitAt n xs of
   (a, b : bs) -> f b <&> \b' -> Just $ a ++ [b'] ++ bs
   _ -> pure Nothing
 
+adjustAtA' :: Applicative f => Int -> (a -> f (a, z)) -> [a] -> f (Maybe ([a], z))
+adjustAtA' n f xs = case splitAt n xs of
+  (a, b : bs) -> f b <&> \(b', z) -> Just (a ++ [b'] ++ bs, z)
+  _ -> pure Nothing
+
 -- | Adjust the first element of the list which satisfies the
 -- predicate. Returns 'Nothing' if there is no such element.
 findAndAdjust :: (a -> Bool) -> (a -> a) -> [a] -> Maybe [a]
@@ -143,6 +150,11 @@ findAndAdjustA p f = \case
   [] -> pure Nothing
   x : xs -> if p x then Just . (: xs) <$> f x else (x :) <<$>> findAndAdjustA p f xs
 
+findAndAdjustA' :: Applicative m => (a -> Bool) -> (a -> m (a, z)) -> [a] -> m (Maybe ([a], z))
+findAndAdjustA' p f = \case
+  [] -> pure Nothing
+  x : xs -> if p x then (\(x', z) -> Just . (,z) . (: xs) $ x') <$> f x else first (x :) <<$>> findAndAdjustA' p f xs
+
 -- | Change the type of an error.
 modifyError :: MonadError e' m => (e -> e') -> ExceptT e m a -> m a
 modifyError f = runExceptT >=> either (throwError . f) pure

primer/src/Primer/API.hs

@@ -127,6 +127,7 @@ import Primer.App (
   newApp,
   progAllDefs,
   progAllTypeDefs,
+  progAllTypeDefsMeta,
   progCxt,
   progImports,
   progLog,
@@ -157,6 +158,7 @@ import Primer.Core (
   TyVarName,
   Type,
   Type' (..),
+  TypeMeta,
   ValConName,
   getID,
   unLocalName,
@@ -218,7 +220,7 @@ import Primer.Log (
 import Primer.Module (moduleDefsQualified, moduleName, moduleTypesQualifiedMeta)
 import Primer.Name qualified as Name
 import Primer.Primitives (primDefType)
-import Primer.TypeDef (ASTTypeDef (..), typeDefNameHints, typeDefParameters)
+import Primer.TypeDef (ASTTypeDef (..), forgetTypeDefMetadata, typeDefNameHints, typeDefParameters)
 import Primer.TypeDef qualified as TypeDef
 import StmContainers.Map qualified as StmMap
 
@@ -1077,23 +1079,23 @@ availableActions ::
 availableActions = curry3 $ logAPI (noError AvailableActions) $ \(sid, level, selection) -> do
   prog <- getProgram sid
   let allDefs = progAllDefs prog
-      allTypeDefs = progAllTypeDefs prog
+      allTypeDefs = progAllTypeDefsMeta prog
   case selection of
     SelectionDef sel -> do
       (editable, ASTDef{astDefType = type_, astDefExpr = expr}) <- findASTDef allDefs sel.def
       pure $ case sel.node of
         Nothing -> Available.forDef (snd <$> allDefs) level editable sel.def
         Just NodeSelection{..} -> case nodeType of
           SigNode -> Available.forSig level editable type_ meta
-          BodyNode -> Available.forBody (snd <$> allTypeDefs) level editable expr meta
+          BodyNode -> Available.forBody (forgetTypeDefMetadata . snd <$> allTypeDefs) level editable expr meta
     SelectionTypeDef sel -> do
-      (editable, _def) <- findASTTypeDef allTypeDefs sel.def
+      (editable, def) <- findASTTypeDef allTypeDefs sel.def
       pure $ case sel.node of
         Nothing -> Available.forTypeDef level editable
         Just (TypeDefParamNodeSelection _) -> Available.forTypeDefParamNode level editable
         Just (TypeDefConsNodeSelection s) -> case s.field of
           Nothing -> Available.forTypeDefConsNode level editable
-          Just _ -> Available.forTypeDefConsFieldNode level editable
+          Just field -> Available.forTypeDefConsFieldNode level editable def s.con field.index field.meta
 
 actionOptions ::
   (MonadIO m, MonadThrow m, MonadAPILog l m) =>
@@ -1117,16 +1119,16 @@ findASTDef allDefs def = case allDefs Map.!? def of
   Just (_, Def.DefPrim _) -> throwM $ UnexpectedPrimDef def
   Just (editable, Def.DefAST d) -> pure (editable, d)
 
-findASTTypeDef :: MonadThrow m => Map TyConName (Editable, TypeDef.TypeDef ()) -> TyConName -> m (Editable, ASTTypeDef ())
+findASTTypeDef :: MonadThrow m => Map TyConName (Editable, TypeDef.TypeDef a) -> TyConName -> m (Editable, ASTTypeDef a)
 findASTTypeDef allTypeDefs def = case allTypeDefs Map.!? def of
   Nothing -> throwM $ UnknownTypeDef def
   Just (_, TypeDef.TypeDefPrim _) -> throwM $ UnexpectedPrimTypeDef def
   Just (editable, TypeDef.TypeDefAST d) -> pure (editable, d)
 
-findASTTypeOrTermDef :: MonadThrow f => App.Prog -> Selection' a -> f (Editable, Either (ASTTypeDef ()) ASTDef)
+findASTTypeOrTermDef :: MonadThrow f => App.Prog -> Selection -> f (Editable, Either (ASTTypeDef TypeMeta) ASTDef)
 findASTTypeOrTermDef prog = \case
   App.SelectionTypeDef sel ->
-    Left <<$>> findASTTypeDef (progAllTypeDefs prog) sel.def
+    Left <<$>> findASTTypeDef (progAllTypeDefsMeta prog) sel.def
   App.SelectionDef sel ->
     Right <<$>> findASTDef (progAllDefs prog) sel.def
 

primer/src/Primer/Action.hs

@@ -1,4 +1,6 @@
+{-# LANGUAGE BlockArguments #-}
 {-# LANGUAGE GADTs #-}
+{-# LANGUAGE OverloadedLabels #-}
 {-# LANGUAGE OverloadedRecordDot #-}
 
 module Primer.Action (
@@ -17,6 +19,7 @@ module Primer.Action (
   uniquifyDefName,
   toProgActionInput,
   toProgActionNoInput,
+  applyActionsToField,
 ) where
 
 import Foreword hiding (mod)
@@ -27,10 +30,11 @@ import Data.Bifunctor.Swap qualified as Swap
 import Data.Generics.Product (typed)
 import Data.List (findIndex)
 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 Optics (set, (%), (?~), (^.), (^?), _Just)
+import Optics (over, set, (%), (?~), (^.), (^?), _Just)
 import Primer.Action.Actions (Action (..), Movement (..), QualifiedText)
 import Primer.Action.Available qualified as Available
 import Primer.Action.Errors (ActionError (..))
@@ -40,6 +44,7 @@ import Primer.App.Base (
   NodeSelection (..),
   NodeType (..),
   Selection' (..),
+  TypeDefConsFieldSelection (..),
   TypeDefConsSelection (..),
   TypeDefNodeSelection (..),
   TypeDefSelection (..),
@@ -60,6 +65,7 @@ import Primer.Core (
   Type' (..),
   TypeCache (..),
   TypeCacheBoth (..),
+  TypeMeta,
   ValConName,
   baseName,
   bindName,
@@ -102,7 +108,7 @@ import Primer.Def (
   Def (..),
   DefMap,
  )
-import Primer.Module (Module, insertDef)
+import Primer.Module (Module (moduleTypes), insertDef)
 import Primer.Name (Name, NameCounter, unName, unsafeMkName)
 import Primer.Name.Fresh (
   isFresh,
@@ -238,6 +244,49 @@ applyActionsToTypeSig smartHoles imports (mod, mods) (defName, def) actions =
             -- In this case we just refocus on the top of the type.
             z -> maybe unwrapError pure (focusType (unfocusLoc z))
 
+applyActionsToField ::
+  (MonadFresh ID m, MonadFresh NameCounter m) =>
+  SmartHoles ->
+  [Module] ->
+  (Module, [Module]) ->
+  (Name, ValConName, Int, ASTTypeDef TypeMeta) ->
+  [Action] ->
+  m (Either ActionError ([Module], TypeZ))
+applyActionsToField smartHoles imports (mod, mods) (tyName, conName', index, tyDef) actions =
+  runReaderT
+    go
+    (buildTypingContextFromModules (mod : mods <> imports) smartHoles)
+    & runExceptT
+  where
+    go :: ActionM m => m ([Module], TypeZ)
+    go = do
+      (valCons, zt) <-
+        (maybe (throwError $ InternalFailure "applyActionsToField: con name not found") pure =<<) $
+          flip (findAndAdjustA' ((== conName') . valConName)) (astTypeDefConstructors tyDef) \(ValCon _ ts) -> do
+            (t, zt) <-
+              maybe (throwError $ InternalFailure "applyActionsToField: con field index out of bounds") pure
+                =<< flip (adjustAtA' index) ts \fieldType -> do
+                  zt <- withWrappedType fieldType \zt ->
+                    foldlM (\l -> local addParamsToCxt . flip applyActionAndSynth l) (InType zt) actions
+                  pure (target (top zt), zt)
+            pure (ValCon conName' t, zt)
+      let mod' = mod{moduleTypes = insert tyName (TypeDefAST tyDef{astTypeDefConstructors = valCons}) $ moduleTypes mod}
+      (,zt) <$> checkEverything smartHoles (CheckEverything{trusted = imports, toCheck = mod' : mods})
+    addParamsToCxt :: TC.Cxt -> TC.Cxt
+    addParamsToCxt = over #localCxt (<> Map.fromList (map (bimap unLocalName TC.K) $ astTypeDefParameters tyDef))
+    withWrappedType :: ActionM m => Type -> (TypeZ -> m Loc) -> m TypeZ
+    withWrappedType ty f = do
+      wrappedType <- ann emptyHole (pure ty)
+      let unwrapError = throwError $ InternalFailure "applyActionsToField: failed to unwrap type"
+          wrapError = throwError $ InternalFailure "applyActionsToField: failed to wrap type"
+          focusedType = focusType $ focus wrappedType
+      case focusedType of
+        Nothing -> wrapError
+        Just wrappedTy ->
+          f wrappedTy >>= \case
+            InType zt -> pure zt
+            z -> maybe unwrapError pure (focusType (unfocusLoc z))
+
 data Refocus = Refocus
   { pre :: Loc
   , post :: Expr
@@ -866,7 +915,7 @@ renameForall b zt = case target zt of
 -- | Convert a high-level 'Available.NoInputAction' to a concrete sequence of 'ProgAction's.
 toProgActionNoInput ::
   DefMap ->
-  Either (ASTTypeDef ()) ASTDef ->
+  Either (ASTTypeDef a) ASTDef ->
   Selection' ID ->
   Available.NoInputAction ->
   Either ActionError [ProgAction]
@@ -958,15 +1007,21 @@ toProgActionNoInput defs def0 sel0 = \case
       typeNodeSel >>= \case
         (s0, TypeDefConsNodeSelection s) -> pure (s0, s)
         _ -> Left NeedTypeDefConsSelection
+    conFieldSel = do
+      (ty, s) <- conSel
+      maybe (Left NeedTypeDefConsFieldSelection) (pure . (ty,s.con,)) s.field
     toProgAction actions = do
-      sel <- termSel
-      toProg' actions sel.def <$> maybeToEither NoNodeSelection sel.node
+      case sel0 of
+        SelectionDef sel -> toProg' actions sel.def <$> maybeToEither NoNodeSelection sel.node
+        SelectionTypeDef _ -> do
+          (t, c, f) <- conFieldSel
+          pure [ConFieldAction t c f.index $ SetCursor f.meta : actions]
     termDef = first (const NeedTermDef) def0
 
 -- | Convert a high-level 'Available.InputAction', and associated 'Available.Option',
 -- to a concrete sequence of 'ProgAction's.
 toProgActionInput ::
-  Either (ASTTypeDef ()) ASTDef ->
+  Either (ASTTypeDef a) ASTDef ->
   Selection' ID ->
   Available.Option ->
   Available.InputAction ->
@@ -1079,9 +1134,15 @@ toProgActionInput def0 sel0 opt0 = \case
     optGlobal = case opt0.context of
       Nothing -> Left $ NeedLocal opt0
       Just q -> pure (q, opt0.option)
+    conFieldSel = do
+      (ty, s) <- conSel
+      maybe (Left NeedTypeDefConsFieldSelection) (pure . (ty,s.con,)) s.field
     toProg actions = do
-      sel <- termSel
-      toProg' actions sel.def <$> maybeToEither NoNodeSelection sel.node
+      case sel0 of
+        SelectionDef sel -> toProg' actions sel.def <$> maybeToEither NoNodeSelection sel.node
+        SelectionTypeDef _ -> do
+          (t, c, f) <- conFieldSel
+          pure [ConFieldAction t c f.index $ SetCursor f.meta : actions]
     offerRefined = do
       id <- nodeID
       def <- termDef

primer/src/Primer/Action/Available.hs

@@ -43,6 +43,11 @@ import Primer.App.Base (
   NodeSelection (..),
   NodeType (..),
   Selection' (..),
+  TypeDefConsFieldSelection (..),
+  TypeDefConsSelection (..),
+  TypeDefNodeSelection (..),
+  TypeDefSelection (..),
+  getTypeDefConFieldType,
  )
 import Primer.Core (
   Expr,
@@ -53,6 +58,8 @@ import Primer.Core (
   ModuleName (unModuleName),
   Type,
   Type' (..),
+  TypeMeta,
+  ValConName,
   getID,
   unLocalName,
   _bindMeta,
@@ -98,6 +105,7 @@ import Primer.Zipper (
   focusOn,
   focusOnTy,
   locToEither,
+  target,
  )
 
 -- | An offered action.
@@ -337,9 +345,15 @@ forTypeDefConsNode l Editable =
 forTypeDefConsFieldNode ::
   Level ->
   Editable ->
+  ASTTypeDef TypeMeta ->
+  ValConName ->
+  Int ->
+  ID ->
   [Action]
-forTypeDefConsFieldNode _ NonEditable = mempty
-forTypeDefConsFieldNode l Editable = sortByPriority l []
+forTypeDefConsFieldNode _ NonEditable _ _ _ _ = mempty
+forTypeDefConsFieldNode l Editable def con index id =
+  maybe mempty (sortByPriority l . forType l) $
+    findType id =<< getTypeDefConFieldType def con index
 
 -- | An input for an 'InputAction'.
 data Option = Option
@@ -375,7 +389,7 @@ options ::
   DefMap ->
   Cxt ->
   Level ->
-  Either (ASTTypeDef ()) ASTDef ->
+  Either (ASTTypeDef TypeMeta) ASTDef ->
   Selection' ID ->
   InputAction ->
   -- | Returns 'Nothing' if an ID was required but not passed, passed but not found in the tree,
@@ -445,9 +459,6 @@ options typeDefs defs cxt level def0 sel0 = \case
   AddCon ->
     pure $ freeVar []
   where
-    defSel = case sel0 of
-      SelectionDef s -> pure s
-      SelectionTypeDef _ -> Nothing
     freeVar opts = Options{opts, free = FreeVarName}
     noFree opts = Options{opts, free = FreeNone}
     localOpt = flip Option Nothing . unName
@@ -461,30 +472,47 @@ options typeDefs defs cxt level def0 sel0 = \case
       pure $
         (first (localOpt . unLocalName) <$> locals)
           <> (first globalOpt <$> globals)
-    findNode = do
-      sel <- defSel
-      s <- sel.node
-      def <- eitherToMaybe def0
-      case s.nodeType of
-        BodyNode -> fst <$> findNodeWithParent s.meta (astDefExpr def)
-        SigNode -> TypeNode <$> findType s.meta (astDefType def)
-    genNames typeOrKind = do
-      sel <- defSel
-      z <- focusNode =<< sel.node
-      pure $ map localOpt $ flip runReader cxt $ case z of
-        Left zE -> generateNameExpr typeOrKind zE
-        Right zT -> generateNameTy typeOrKind zT
-    varsInScope = do
-      sel <- defSel
-      nodeSel <- sel.node
-      focusNode nodeSel <&> \case
-        Left zE -> variablesInScopeExpr defs zE
-        Right zT -> (variablesInScopeTy zT, [], [])
+    findNode = case sel0 of
+      SelectionDef sel -> do
+        nodeSel <- sel.node
+        def <- eitherToMaybe def0
+        case nodeSel.nodeType of
+          BodyNode -> fst <$> findNodeWithParent nodeSel.meta (astDefExpr def)
+          SigNode -> TypeNode <$> findType nodeSel.meta (astDefType def)
+      SelectionTypeDef sel -> do
+        (_, zT) <- conField sel
+        pure $ TypeNode $ target zT
+    genNames typeOrKind =
+      map localOpt . flip runReader cxt <$> case sel0 of
+        SelectionDef sel -> do
+          z <- focusNode =<< sel.node
+          pure $ case z of
+            Left zE -> generateNameExpr typeOrKind zE
+            Right zT -> generateNameTy typeOrKind zT
+        SelectionTypeDef sel -> do
+          (_, zT) <- conField sel
+          pure $ generateNameTy typeOrKind zT
+    varsInScope = case sel0 of
+      SelectionDef sel -> do
+        nodeSel <- sel.node
+        focusNode nodeSel <&> \case
+          Left zE -> variablesInScopeExpr defs zE
+          Right zT -> (variablesInScopeTy zT, [], [])
+      SelectionTypeDef sel -> do
+        (def, zT) <- conField sel
+        pure (astTypeDefParameters def <> variablesInScopeTy zT, [], [])
     focusNode nodeSel = do
       def <- eitherToMaybe def0
       case nodeSel.nodeType of
         BodyNode -> Left . locToEither <$> focusOn nodeSel.meta (astDefExpr def)
         SigNode -> fmap Right $ focusOnTy nodeSel.meta $ astDefType def
+    conField sel = do
+      (con, field) <- case sel of
+        TypeDefSelection _ (Just (TypeDefConsNodeSelection (TypeDefConsSelection con (Just field)))) ->
+          Just (con, field)
+        _ -> Nothing
+      def <- either Just (const Nothing) def0
+      map (def,) $ focusOnTy field.meta =<< getTypeDefConFieldType def con field.index
     -- Extract the source of the function type we were checked at
     -- i.e. the type that a lambda-bound variable would have here
     lamVarTy = \case