feat: extensible slicing

Make syntax slicing extensible so that custom slicers can be written
for various productions that aren't handled well by the default
heuristics.

src/compat/SubVerso/Compat.lean

@@ -1,4 +1,4 @@
-import Lean
+import Lean.Elab
 
 open Lean Elab Term
 

src/examples/SubVerso/Examples/Slice.lean

@@ -6,13 +6,15 @@ Author: David Thrane Christiansen
 import Lean.Data.Parsec
 import Lean.Data.Position
 import Lean.Syntax
-
 import Lean.Elab.Command
+import SubVerso.Examples.Slice.Attribute
 
 open Lean
 
 namespace SubVerso.Examples.Slice
 
+open Attribute
+
 private def takeComments (ws : Substring) : Array (SourceInfo × Substring) := Id.run do
   let mut iter : String.Iterator := ws.toIterator
   let mut leading := iter
@@ -118,27 +120,43 @@ where
     | _ => #[]
 
 
+/--
+Returns `true` if `stx`'s source span is definitely included in `rng`.
+
+Returns false in the following circumstances:
+ * No range can be found for the syntax
+ * A range is found, but it exceeds rng in one or both directions or does not overlap it
+-/
 private def syntaxIn (stx : Syntax) (rng : String.Range) : Bool :=
   if let some sr := stx.getRange? then
     rng.includes sr
   else
     false
 
-private partial def removeRanges (rngs : Array String.Range) (stx : Syntax) : Syntax :=
+private partial def removeRanges (env : Environment) (rngs : Array String.Range) (stx : Syntax) : Syntax :=
   match stx with
   | .atom .. | .ident .. | .missing => stx
-  | .node info kind subs =>
+  | .node info kind subs => Id.run do
     if rngs.size > 0 then
-      let newSubs := subs.filter (fun s => !rngs.any (syntaxIn s)) |>.map (removeRanges rngs)
+      let slicers ← slicersInEnv env kind
+
+      for s in slicers do
+        if let some stx' := s (removeRanges env rngs) (.node info kind subs) rngs then
+          return stx'
+
+      -- Fallback if no special-purpose slicers apply
+      let newSubs ← subs.filter (fun s => !rngs.any (syntaxIn s)) |>.map (removeRanges env rngs)
+
       -- If there were atoms in the original tree and only one non-atom subtree remains, replace the surrounding node
       -- This catches situations like
       --    x /- !! begin foo  -/ + 3 /- !! end foo -/
       -- which should become just x.
-      if h : subs.any (· matches .atom ..) ∧ subs.size > 1 ∧ newSubs.size = 1 then
-        newSubs[0]'(by cases h; simp_arith [*])
+      if h : subs.any (· matches .atom ..) ∧ subs.size > 1 ∧ newSubs.size = 1 ∧ newSubs[0]? matches some (Syntax.node ..) then
+        pure <| newSubs[0]'(by cases h; simp_arith [*])
       else
-        .node info kind newSubs
-    else stx
+        pure <| .node info kind newSubs
+    else pure stx
+
 
 private def getSlices (slices : Array SliceCommand) : Except String (HashMap String (Array String.Range)) := do
   let mut opened : HashMap String (String.Pos) := {}
@@ -168,6 +186,9 @@ where
   getRange : Substring → String.Range
     | ⟨_, b, e⟩ => ⟨b, e⟩
   remFromSubstring (str : Substring) : Substring := Id.run do
+    -- Not just an optimization - some parsers use empty strings that aren't build as substrings
+    -- of the original source
+    if str.str.isEmpty then return str
     if str.str != fileMap.source then panic! "Mismatched strings"
     let mut found := false
     let mut start := str.startPos
@@ -201,19 +222,47 @@ structure Slices where
   residual : Syntax
   sliced : HashMap String Syntax
 
-def sliceSyntax [Monad m] [MonadError m] [MonadFileMap m] (stx : Syntax) : m Slices := do
+def sliceSyntax [Monad m] [MonadEnv m] [MonadError m] [MonadFileMap m] (stx : Syntax) : m Slices := do
   let commands := getComments stx |>.filterMap (fun (_, _, cmd?) => cmd?)
   let ss ←
     match getSlices commands with
     | .ok ss => pure ss
     | .error m => throwError m
 
   let stx' := removeSliceComments commands (← getFileMap) stx
-
+  let env ← getEnv
   let mut sliced : HashMap String Syntax := {}
   for (s, _) in ss.toArray do
     let rem := ss.erase s |>.toList |>.bind (·.snd.toList) |>.toArray
-    sliced := sliced.insert s (removeRanges rem stx')
-  let residual := removeRanges (ss.toList.bind (·.snd.toList) |>.toArray) stx'
+    sliced := sliced.insert s (removeRanges env rem stx')
+  let residual := removeRanges env (ss.toList.bind (·.snd.toList) |>.toArray) stx'
 
   return {original := stx, residual, sliced}
+
+namespace Syntax
+
+
+@[slicer null]
+def null : Slicer := fun go stx rngs => do
+  let .node si k subs := stx
+    | none
+  let newSubs := subs.filter (fun s => !rngs.any (syntaxIn s)) |>.map go
+  -- Don't apply the heuristic for null nodes
+  return .node si k newSubs
+
+@[slicer Lean.Parser.Tactic.tacticSeq1Indented]
+def tacticSeq1Indented : Slicer := fun go stx rngs => do
+  let .node si k #[.node si' k' subs] := stx
+    | none
+  let mut out : Array Syntax := #[]
+  let mut i := 0
+  while h : i < subs.size do
+    if !rngs.any (syntaxIn subs[i]) then
+      out := out.push (go subs[i])
+      if h' : i + 1 < subs.size then
+        out := out.push subs[i+1]
+      else
+        out := out.push (.node SourceInfo.none nullKind #[])
+    i := i + 2
+  -- Pop the resulting array because the loop pushes an extraneous null separator node onto the end
+  return .node si k #[.node si' k' out.pop]

src/examples/SubVerso/Examples/Slice/Attribute.lean

@@ -0,0 +1,45 @@
+/-
+Copyright (c) 2024 Lean FRO LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Author: David Thrane Christiansen
+-/
+import Lean.Data.Position
+import Lean.Syntax
+import Lean.KeyedDeclsAttribute
+
+namespace SubVerso.Examples.Slice.Attribute
+
+open Lean
+
+private unsafe def mkSliceExpanderAttrUnsafe (attrName typeName : Name) (descr : String) (attrDeclName : Name) : IO (KeyedDeclsAttribute α) :=
+  KeyedDeclsAttribute.init {
+    name := attrName,
+    descr := descr,
+    valueTypeName := typeName
+  } attrDeclName
+
+@[implemented_by mkSliceExpanderAttrUnsafe]
+private opaque mkSliceExpanderAttributeSafe (attrName typeName : Name) (desc : String) (attrDeclName : Name) : IO (KeyedDeclsAttribute α)
+
+private def mkSliceExpanderAttribute (attrName typeName : Name) (desc : String) (attrDeclName : Name := by exact decl_name%) : IO (KeyedDeclsAttribute α) :=
+  mkSliceExpanderAttributeSafe attrName typeName desc attrDeclName
+
+abbrev Slicer := (Syntax → Syntax) → Syntax → Array String.Range → Option Syntax
+
+initialize slicerAttr : KeyedDeclsAttribute Slicer ←
+  mkSliceExpanderAttribute `slicer ``Slicer "Indicates that this function is used to slice the given syntax kind"
+
+unsafe def slicersForUnsafe [Monad m] [MonadEnv m] (x : Name) : m (List Slicer) := do
+  let expanders := slicerAttr.getEntries (← getEnv) x
+  return expanders.map (·.value)
+
+@[implemented_by slicersForUnsafe]
+opaque slicersFor [Monad m] [MonadEnv m] (x : Name) : m (List Slicer)
+
+
+unsafe def slicersInEnvUnsafe (env : Environment) (x : Name) : List Slicer :=
+  let expanders := slicerAttr.getEntries env x
+  expanders.map (·.value)
+
+@[implemented_by slicersInEnvUnsafe]
+opaque slicersInEnv (env : Environment) (x : Name) : List Slicer