Sort-dependent expansion for Modules

src/haz3lcore/TyDi/TyDiForms.re

@@ -100,13 +100,16 @@ module Delims = {
   let leading = (sort: Sort.t): list(Token.t) =>
     Form.delims
     |> List.map(token => {
-         let (lbl, _) = Form.Expansion.get(token);
-         List.filter_map(
-           (m: Mold.t) =>
-             List.length(lbl) > 1 && token == List.hd(lbl) && m.out == sort
-               ? Some(token ++ leading_expander) : None,
-           Form.Molds.get(lbl),
-         );
+         let (lbl, _) = Form.Expansion.get(sort, token);
+         switch (Form.Molds.try_get(sort, lbl)) {
+         | None => []
+         | Some(molds) =>
+           molds
+           |> List.filter_map((_: Mold.t) =>
+                List.length(lbl) > 1 && token == List.hd(lbl)
+                  ? Some(token ++ leading_expander) : None
+              )
+         };
        })
     |> List.flatten
     |> List.sort_uniq(compare);
@@ -151,12 +154,14 @@ module Delims = {
   let const_mono = (sort: Sort.t): list(Token.t) =>
     Token.const_mono_delims
     |> List.map(token => {
-         List.filter_map(
-           (m: Mold.t) =>
-             m.out == sort && List.mem(token, Token.const_mono_delims)
-               ? Some(token) : None,
-           Form.Molds.get([token]),
-         )
+         switch (Form.Molds.try_get(sort, [token])) {
+         | None => []
+         | Some(molds) =>
+           molds
+           |> List.filter_map((_: Mold.t) =>
+                List.mem(token, Token.const_mono_delims) ? Some(token) : None
+              )
+         }
        })
     |> List.flatten
     |> List.sort_uniq(compare);

src/haz3lcore/lang/Form.re

@@ -31,6 +31,10 @@ type expansion =
 [@deriving (show({with_path: false}), sexp, yojson)]
 type expansions = list((Token.t, (Label.t, Direction.t)));
 
+/* Sort-aware expansions include the form's sort for filtering */
+[@deriving (show({with_path: false}), sexp, yojson)]
+type sorted_expansions = list((Token.t, Sort.t, Label.t, Direction.t));
+
 /* A label, a mold, and expansion behavior together determine a form. */
 [@deriving (show({with_path: false}), sexp, yojson)]
 type t = {
@@ -196,6 +200,10 @@ type compound_form =
   | Use
   // TRIPLE DELIMITERS
   | Let
+  | TEST_Let //TODO(andrew): rm
+  | TEST_TypeAlias //TODO(andrew): rm
+  | TEST_Seq
+  | TEST_Curly
   | Theorem
   | TypeAlias
   | If;
@@ -292,6 +300,10 @@ let get: compound_form => t =
   | ProofOf => mk_op_c(L, ["proof_of", "end"], Typ, [Exp])
   // TRIPLE DELIMITERS
   | Let => mk_pre_c(L, ["let", "=", "in"], P.let_, Exp, [Pat, Exp])
+  | TEST_Let => mk_pre_c(L, ["let", "="], P.let_, Pat, [Pat])
+  | TEST_TypeAlias => mk_pre_c(L, ["type", "="], P.let_, Pat, [TPat])
+  | TEST_Seq => mk_infix(";", Pat, P.semi)
+  | TEST_Curly => mk_op_c(LT, ["{", "}"], Pat, [Pat])
   | TypeAlias => mk_pre_c(L, ["type", "=", "in"], P.let_, Exp, [TPat, Typ])
   | If => mk_pre_c(L, ["if", "then", "else"], P.if_, Exp, [Exp, Exp])
   | HintedTest => mk_op_c(L, ["hint", "test", "end"], Exp, [Exp, Exp]);
@@ -423,26 +435,45 @@ module Molds = {
   let compound = (label: Label.t): option(list(Mold.t)) =>
     List.assoc_opt(label, compounds);
 
-  let get = (label: Label.t): list(Mold.t) =>
+  /* Base: get molds from form definitions without sort filtering */
+  let get_base = (label: Label.t): list(Mold.t) =>
     switch (label, compound(label)) {
     | ([t], Some(molds)) when atomic(t) != [] => atomic(t) @ molds
     | ([t], None) when atomic(t) != [] => atomic(t)
     | (_, Some(molds)) => molds
-    /* For tokens which are not assigned molds by the language definition,
-       we assign a default 'Any' mold, which is either convex or concave.
-       This is half-ass at the moment as we don't have a rigorous lexing
-       policy, but is somewhat load-bearing in that remolding as one is
-       typing in a multi-character operator can cause jank, which is
-       alleviated if we correctly guess that it will become an operator. */
-    | ([t], None)
-        when Token.is_potential_operator(t) && !Token.is_potential_operand(t) => [
-        Mold.mk_bin(Precedence.max, Any, []),
-      ]
-    | (_, None) => [Mold.mk_op(Any, [])]
+    | _ => []
+    };
+
+  /* Strict: filter by sort, returns None if no match.
+     Used by remolding where we need accurate sort filtering. */
+  let try_get = (sort: Sort.t, label: Label.t): option(list(Mold.t)) => {
+    let molds = get_base(label);
+    let filtered = molds |> List.filter((m: Mold.t) => m.out == sort);
+    filtered == [] ? None : Some(filtered);
+  };
+
+  /* Get mold for insertion: permissive sort filtering with fallback
+     to Any-sorted default molds for undefined tokens. */
+  let get = (sort: Sort.t, label: Label.t): Mold.t =>
+    switch (try_get(sort, label)) {
+    | Some(molds) =>
+      assert(molds != []);
+      List.hd(molds);
+    | None =>
+      /* Fallback: create Any-sorted default mold. This handles tokens
+         not assigned molds by the language definition. */
+      switch (label) {
+      | [t]
+          when
+            Token.is_potential_operator(t) && !Token.is_potential_operand(t) =>
+        Mold.mk_bin(Precedence.max, Any, [])
+      | _ => Mold.mk_op(Any, [])
+      }
     };
 };
 
 module Expansion = {
+  /* Sort-agnostic expansion info (for backward compatibility) */
   let expanding_of = ({expansion, label, _}: t): option(expansions) =>
     switch (expansion, label) {
     | (L, [hd, ..._]) => Some([(hd, (label, Direction.Left))])
@@ -451,19 +482,75 @@ module Expansion = {
     | _ => None
     };
 
+  /* Sort-aware expansion info - uses nib sorts for context matching.
+     Leading delimiters use left nib sort (the context you're in when typing).
+     Trailing delimiters use right nib sort.
+
+     Note: This uses nib sort rather than mold.out because the nib sort
+     reflects what context you're typing in, not what the form produces.
+     For example, Rule ["|", "=>"] has out=Rul but left nib=Exp, since
+     you type | after an expression (the previous rule body).
+
+     Limitation: Ascriptions (expr : Type) have Typ right nib even though
+     they produce Exp. This causes issues for forms like | that can follow
+     ascribed expressions. See Insert.re for the special case handling. */
+  let sorted_expanding_of =
+      ({expansion, label, mold}: t): option(sorted_expansions) => {
+    let (l_nib, r_nib) = mold.nibs;
+    switch (expansion, label) {
+    | (L, [hd, ..._]) => Some([(hd, l_nib.sort, label, Direction.Left)])
+    | (LT, [hd, ..._]) =>
+      Some([
+        (hd, l_nib.sort, label, Left),
+        (ListUtil.last(label), r_nib.sort, label, Right),
+      ])
+    | _ => None
+    };
+  };
+
+  /* Sort-agnostic expansions (kept for is_leading) */
   let expansions: expansions =
     List.filter_map(((_, form: t)) => expanding_of(form), forms)
     |> List.flatten
     |> List.sort_uniq(compare);
 
-  let get = (t: Token.t): (Label.t, Direction.t) =>
-    switch (List.assoc_opt(t, expansions)) {
-    | Some(expansion) => expansion
-    | None => ([t], Right)
+  /* Sort-aware expansions */
+  let sorted_expansions: sorted_expansions =
+    List.filter_map(((_, form: t)) => sorted_expanding_of(form), forms)
+    |> List.flatten;
+
+  /* Get expansion for a token in a specific sort context.
+     Returns monotile if no expansion exists for this sort.
+     Exception: Rul context is permissive - falls back to any expansion.
+     This is because Rul (case rules) contains Exp/Pat operands but has no
+     direct forms for things like parens. Other sorts remain strict. */
+  let get = (sort: Sort.t, t: Token.t): (Label.t, Direction.t) => {
+    let matching =
+      sorted_expansions
+      |> List.find_opt(((tok, s, _, _)) => tok == t && s == sort);
+    switch (matching) {
+    | Some((_, _, lbl, dir)) => (lbl, dir)
+    | None =>
+      switch (sort) {
+      | Rul =>
+        /* Rul context: fall back to any expansion since rules contain
+           Exp/Pat operands but have no direct operand forms. */
+        let any_match =
+          sorted_expansions |> List.find_opt(((tok, _, _, _)) => tok == t);
+        switch (any_match) {
+        | Some((_, _, lbl, dir)) => (lbl, dir)
+        | None => ([t], Right)
+        };
+      | _ => ([t], Right)
+      }
     };
+  };
 
-  let will = kw => List.length(get(kw) |> fst) > 1;
+  /* Check if token would expand in ANY sort (sort-agnostic) */
+  let will = (t: Token.t): bool =>
+    List.exists(((tok, _, _, _)) => tok == t, sorted_expansions);
 
+  /* Check if token is a leading delimiter in ANY sort (sort-agnostic) */
   let is_leading = (t: Token.t): bool =>
     switch (List.assoc_opt(t, expansions)) {
     | Some((_, Left)) => true

src/haz3lcore/tiles/Segment.re

@@ -107,21 +107,24 @@ let rec remold = (~shape=Nib.Shape.concave(), seg: t, s: Sort.t) =>
 and remold_tile = (s: Sort.t, shape, t: Tile.t): option(Tile.t) => {
   open OptUtil.Syntax;
   let+ remolded =
-    Form.Molds.get(t.label)
-    |> List.filter((m: Mold.t) => m.out == s)
-    |> List.map(mold =>
-         {
-           ...t,
-           mold,
-         }
-       )
-    |> (
-      fun
-      | [_] as ts => ts
-      | ts =>
-        ts |> List.filter(t => Nib.Shape.fits(shape, fst(Tile.shapes(t))))
-    )
-    |> ListUtil.hd_opt;
+    switch (Form.Molds.try_get(s, t.label)) {
+    | None => None
+    | Some(molds) =>
+      molds
+      |> List.map(mold =>
+           {
+             ...t,
+             mold,
+           }
+         )
+      |> (
+        fun
+        | [_] as ts => ts
+        | ts =>
+          ts |> List.filter(t => Nib.Shape.fits(shape, fst(Tile.shapes(t))))
+      )
+      |> ListUtil.hd_opt
+    };
   let children =
     List.fold_right(
       ((l, child, r), children) => {

src/haz3lcore/zipper/Ancestor.re

@@ -27,10 +27,6 @@ let nibs = (a: t) => {
   let (_, r) = Mold.nibs(~index=r_shard(a), a.mold);
   (l, r);
 };
-let shapes = a => {
-  let (l, r) = nibs(a);
-  (l.shape, r.shape);
-};
 
 let zip = (child: Segment.t, {id, label, mold, shards, children}: t): Tile.t => {
   id,
@@ -40,22 +36,6 @@ let zip = (child: Segment.t, {id, label, mold, shards, children}: t): Tile.t =>
   children: fst(children) @ [child, ...snd(children)],
 };
 
-let sorted_children = (a: t) => {
-  let n = List.length(fst(a.children));
-  let t = zip(Segment.empty, a);
-  let (l, _, r) = ListUtil.split_nth(n, Tile.sorted_children(t));
-  (l, r);
-};
-
-let remold = (a: t): list(t) =>
-  Form.Molds.get(a.label)
-  |> List.map(mold =>
-       {
-         ...a,
-         mold,
-       }
-     );
-
 let sort = (a: t): Sort.t => {
   let (pre, suf) = a.shards;
   switch (ListUtil.split_last_opt(pre), suf) {
@@ -88,29 +68,6 @@ let missing_middle_shards = (a: t): list(Tile.t) => {
   Tile.split_shards(a.id, a.label, a.mold, ls);
 };
 
-let missing_shards = (a: t): list(Tile.t) => {
-  let (shards_l, shards_r) = a.shards;
-  let shards = shards_l @ shards_r;
-  let missing =
-    List.filter(
-      i => !List.mem(i, shards),
-      List.init(List.length(a.label), Fun.id),
-    );
-  Tile.split_shards(a.id, a.label, a.mold, missing);
-};
-
-let container_shards = (a: t): (Piece.t, Piece.t) => {
-  let (shards_l, shards_r) =
-    a.shards
-    |> TupleUtil.map2(Tile.split_shards(a.id, a.label, a.mold))
-    |> TupleUtil.map2(List.map(Tile.to_piece));
-  let l =
-    ListUtil.last_opt(shards_l) |> OptUtil.get_or_raise(Empty_shard_affix);
-  let r =
-    ListUtil.hd_opt(shards_r) |> OptUtil.get_or_raise(Empty_shard_affix);
-  (l, r);
-};
-
 let reassemble = (match_l: Aba.t(Tile.t, Segment.t) as 'm, match_r: 'm): t => {
   // TODO(d) bit hacky, need to do a flip/orientation pass
   // let match_l = Aba.map_b(Segment.rev, match_l);

src/haz3lcore/zipper/Relatives.re

@@ -80,6 +80,29 @@ let delete_parent = ({siblings, ancestors}: t): t => {
   };
 };
 
+/* The sort at the current insertion point, accounting for
+ * infix operators with heterogeneous child sorts (e.g. type
+ * annotation ':' in patterns). This looks at the right nib
+ * of the last tile to the left, which determines what sort
+ * should come next - the same logic used by Segment.remold. */
+let sort = ({siblings: (pre, _), ancestors}: t): Sort.t => {
+  let outer_sort = Ancestors.sort(ancestors);
+  let rec find_last_tile =
+    fun
+    | [] => None
+    | [p, ...rest] =>
+      switch (Piece.is_tile(p)) {
+      | Some(t) => Some(t)
+      | None => find_last_tile(rest)
+      };
+  switch (find_last_tile(List.rev(pre))) {
+  | None => outer_sort
+  | Some(t) =>
+    let (_, r) = Tile.nibs(t);
+    r.sort;
+  };
+};
+
 let remold = ({siblings, ancestors}: t): t => {
   let s = Ancestors.sort(ancestors);
   let siblings = Siblings.remold(siblings, s);