feat: syntax slices

Specially-formatted comments allow mutliple versions of syntax to be
defined at once.

This is not exposed in the example serialization API yet, but the plan
is to include it there as well.

Tests.lean

@@ -1,12 +1,58 @@
 import SubVerso.Highlighting
 import SubVerso.Examples
 import SubVerso.Examples.Env
+import SubVerso.Examples.Slice
 import Lean.Data.Json
 import Lean.Data.NameMap
 
 open SubVerso.Examples (loadExamples Example)
 open Lean.FromJson (fromJson?)
 
+
+section
+open Lean Elab Command
+open SubVerso.Examples.Slice
+
+elab "#test_slicing" s:str out:ident c:command : command => do
+  let {original, residual, sliced} ← sliceSyntax c
+
+  let mut log := ""
+
+  log := log ++ s!"Original:\n{Syntax.prettyPrint c}"
+  log := log ++ s!"Residual:\n{Syntax.prettyPrint residual}"
+  for (s, stx) in sliced.toArray do
+    log := log ++ s!"Slice {s}:\n{Syntax.prettyPrint stx}"
+
+  let use := s.getString
+  if let some stx := sliced.find? use then
+    elabCommand stx
+  else
+    logInfo m!"Slice '{use}' not found"
+    elabCommand residual
+  elabCommand <| ← `(def $out := $(quote log))
+
+set_option pp.rawOnError true in
+#test_slicing "foo" sliceLog
+def heya : IO Unit := do
+  IO.print "Hey"
+  -- !! begin foo
+  IO.println " there"
+  /- Bigger comment
+  -/
+  -- !! end foo
+  let x := 33 -- !! begin foo
+  if x > 3 then pure () -- !! end foo
+  if x == 3 /-
+    !! begin foo
+  -/ || true /- !! end foo -/ then pure ()
+  let x := [1, 2/- !! begin foo-/, 3/- !! end foo-/]
+  pure ()
+
+def expectedLog :=
+  "Original:\ndef heya : IO Unit := do\n  IO.print \"Hey\"\n  -- !! begin foo\n  IO.println \" there\"\n  /- Bigger comment\n  -/\n  -- !! end foo\n  let x := 33 -- !! begin foo\n  if x > 3 then pure () -- !! end foo\n  if x == 3 /-\n    !! begin foo\n  -/ || true /- !! end foo -/ then pure ()\n  let x := [1, 2/- !! begin foo-/, 3/- !! end foo-/]\n  pure ( ) Residual:\ndef heya : IO Unit := do\n  IO.print \"Hey\"\n  let x := 33 \n  if x == 3  then pure ()\n  let x := [1, 2]\n  pure ( ) Slice foo:\ndef heya : IO Unit := do\n  IO.print \"Hey\"\n  IO.println \" there\"\n  /- Bigger comment\n  -/\n  let x := 33 \n  if x > 3 then pure () \n  if x == 3  || true  then pure ()\n  let x := [1, 2, 3]\n  pure ( ) "
+
+end
+
 open Lean in
 def exampleArray (examples : NameMap (NameMap α)) : Array α := Id.run do
   let mut exs := #[]
@@ -67,6 +113,13 @@ def checkVersion (expected : String) (examples : NameMap (NameMap Example)) : IO
   pure ()
 
 def main : IO UInt32 := do
+  IO.println "Checking the slice log"
+  if sliceLog != expectedLog then
+    IO.println "Mismatch between expected:"
+    IO.println sliceLog
+    IO.println "and actual:"
+    IO.println expectedLog
+
   IO.println "Checking that the test project generates at least some deserializable JSON"
   cleanupDemo
   let examples ← loadExamples "demo"

src/examples/SubVerso/Examples/Slice.lean

@@ -0,0 +1,219 @@
+/-
+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.Parsec
+import Lean.Data.Position
+import Lean.Syntax
+
+import Lean.Elab.Command
+
+open Lean
+
+namespace SubVerso.Examples.Slice
+
+private def takeComments (ws : Substring) : Array (SourceInfo × Substring) := Id.run do
+  let mut iter : String.Iterator := ws.toIterator
+  let mut leading := iter
+  let mut found := #[]
+  while iter.pos < ws.stopPos do
+    if iter.curr == '-' && iter.hasNext then
+      let mut iter' := iter.next
+      if iter'.curr == '-' && iter'.hasNext then
+        while iter'.pos < ws.stopPos && iter'.curr != '\n' do
+          iter' := iter'.next
+        let l := ⟨ws.str, leading.pos, iter.pos⟩
+        let t := ⟨ws.str, iter'.pos, iter'.pos⟩
+        found := found.push ((.original l iter.pos t iter'.pos), ⟨iter.s, iter.i, iter'.i⟩)
+        iter := iter'
+        continue
+    else if iter.curr == '/' then
+      let mut iter' := iter.next
+      if iter'.curr == '-' && iter'.hasNext then
+        let mut depth := 1
+        while depth > 0 && iter'.pos < ws.stopPos && iter'.hasNext do
+          if iter'.curr == '/' && iter'.hasNext then
+            iter' := iter'.next
+            if iter'.curr == '-' then
+              depth := depth + 1
+              iter' := iter'.next
+          else if iter'.curr == '-' && iter'.hasNext then
+            iter' := iter'.next
+            if iter'.curr == '/' then
+              depth := depth - 1
+              iter' := iter'.next
+          else
+            iter' := iter'.next
+        let l := ⟨ws.str, leading.pos, iter.pos⟩
+        let t := ⟨ws.str, iter'.pos, iter'.pos⟩
+        found := found.push ((.original l iter.pos t iter'.pos), ⟨iter.s, iter.i, iter'.i⟩)
+        iter := iter'
+    else
+      iter := iter.next
+  found
+
+private def commentString (str : Substring) : Substring :=
+  if str.take 2 == "--".toSubstring then str.drop 2
+  else str.drop 2 |>.dropRight 2
+
+private structure SliceCommand where
+  begin : Bool
+  kind : String
+  commentAt : String.Range
+deriving Repr
+
+private def SliceCommand.beginRange := SliceCommand.mk true
+private def SliceCommand.endRange := SliceCommand.mk false
+
+open Parsec in
+private def sliceCommand (range : String.Range) : Parsec SliceCommand := do
+  ws
+  skipString "!!"
+  ws
+  beginCmd <|> endCmd
+where
+  nameChar (c : Char) := c.isAlphanum || c == '_'
+  done : Parsec Unit := fun iter =>
+    if iter.pos ≥ range.stop then .success iter ()
+    else .error iter "Expected finished comment"
+  beginCmd := do
+    skipString "begin"
+    skipChar ' '
+    ws
+    let name ← many1Chars (satisfy nameChar)
+    ws
+    _ ← optional <| skipString "-/"
+    done
+    pure <| .beginRange name range
+  endCmd := do
+    skipString "end"
+    skipChar ' '
+    ws
+    let name ← many1Chars (satisfy nameChar)
+    ws
+    _ ← optional <| skipString "-/"
+    done
+    pure <| .endRange name range
+
+
+private partial def getComments : Syntax → Array (SourceInfo × Substring × Option SliceCommand)
+  | .atom si _str => ws si
+  | .ident si _str _x _ => ws si
+  | .node si _k sub =>
+    sub.foldr (fun a b => getComments a ++ b) (ws si)
+  | .missing => #[]
+where
+  ws
+    | .original leading _ trailing _ =>
+      (takeComments leading ++ takeComments trailing).map fun (si, str) =>
+        let iter := commentString str |>.toIterator
+        match si with
+        | .original _ b _ e =>
+          match sliceCommand ⟨b, e⟩ iter with
+          | .error _iter _msg =>
+            (si, str, none)
+          | .success _ v => (si, str, some v)
+        | _ => (si, str, none)
+    | _ => #[]
+
+
+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 :=
+  match stx with
+  | .atom .. | .ident .. | .missing => stx
+  | .node info kind subs =>
+    if rngs.size > 0 then
+      let newSubs := subs.filter (fun s => !rngs.any (syntaxIn s)) |>.map (removeRanges 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 [*])
+      else
+        .node info kind newSubs
+    else stx
+
+private def getSlices (slices : Array SliceCommand) : Except String (HashMap String (Array String.Range)) := do
+  let mut opened : HashMap String (String.Pos) := {}
+  let mut closed : HashMap String (Array String.Range) := {}
+  for s in slices do
+    match s with
+    | .beginRange name rng =>
+      if opened.contains name then .error s!"Slice '{name}' opened twice"
+      else opened := opened.insert name rng.stop
+    | .endRange name pos =>
+      if let some p := opened.find? name then
+        closed := closed.insert name <| closed.findD name #[] |>.push ⟨p, pos.start⟩
+        opened := opened.erase name
+      else .error s!"Slice '{name}' not open"
+  pure closed
+
+private partial def removeSliceComments (slices : Array SliceCommand) (fileMap : FileMap) : Syntax → Syntax
+  | .atom si str => .atom (remFromSourceInfo si) str
+  | .ident si str n pres => .ident (remFromSourceInfo si) str n pres
+  | .missing => .missing
+  | .node si k sub => .node (remFromSourceInfo si) k (sub.map (removeSliceComments slices fileMap))
+where
+  remFromSourceInfo : SourceInfo → SourceInfo
+    | .original leading b trailing e =>
+      .original (remFromSubstring leading) b (remFromSubstring trailing) e
+    | other => other
+  getRange : Substring → String.Range
+    | ⟨_, b, e⟩ => ⟨b, e⟩
+  remFromSubstring (str : Substring) : Substring := Id.run do
+    if str.str != fileMap.source then panic! "Mismatched strings"
+    let mut found := false
+    let mut start := str.startPos
+    let mut newStr := ""
+    for s in slices do
+      if getRange str |>.includes s.commentAt then
+        found := true
+        -- If the characters on the same line as the start and end of the comment are all whitespace,
+        -- delete whole lines; otherwise, delete just the comment's range
+        let startPos := fileMap.toPosition s.commentAt.start
+        let startLine := fileMap.positions[startPos.line - 1]!
+        let preComment := fileMap.source.extract startLine s.commentAt.start
+
+        let endPos := fileMap.toPosition s.commentAt.stop
+        let endLine := fileMap.positions[endPos.line]!
+        let postComment := fileMap.source.extract s.commentAt.stop endLine
+        let deleteLines := preComment.all Char.isWhitespace && postComment.all Char.isWhitespace
+
+        if deleteLines then
+          newStr := fileMap.source.extract start startLine
+          start := endLine
+        else
+          newStr := fileMap.source.extract start s.commentAt.start
+          start := s.commentAt.stop
+
+    newStr := newStr ++ fileMap.source.extract start str.stopPos
+    if found then newStr.toSubstring else str
+
+structure Slices where
+  original : Syntax
+  residual : Syntax
+  sliced : HashMap String Syntax
+
+def sliceSyntax [Monad 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 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'
+
+  return {original := stx, residual, sliced}