Factor ast.Slice out into its own package

experimental/ast/commas.go

@@ -0,0 +1,70 @@
+// Copyright 2020-2024 Buf Technologies, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package ast
+
+import (
+	"slices"
+
+	"github.com/bufbuild/protocompile/experimental/seq"
+	"github.com/bufbuild/protocompile/experimental/token"
+)
+
+// Commas is like [Slice], but it's for a comma-delimited list of some kind.
+//
+// This makes it easy to work with the list as though it's a slice, while also
+// allowing access to the commas.
+type Commas[T any] interface {
+	seq.Inserter[T]
+
+	// Comma is like [seq.Indexer.At] but returns the comma that follows the nth
+	// element.
+	//
+	// May be [token.Zero], either because it's the last element
+	// (a common situation where there is no comma) or it was added with
+	// Insert() rather than InsertComma().
+	Comma(n int) token.Token
+
+	// AppendComma is like [seq.Append], but includes an explicit comma.
+	AppendComma(value T, comma token.Token)
+
+	// InsertComma is like [seq.Inserter.Insert], but includes an explicit comma.
+	InsertComma(n int, value T, comma token.Token)
+}
+
+type withComma[T any] struct {
+	Value T
+	Comma token.ID
+}
+
+type commas[T, E any] struct {
+	seq.SliceInserter[T, withComma[E]]
+	ctx Context
+}
+
+func (c commas[T, _]) Comma(n int) token.Token {
+	return (*c.SliceInserter.Slice)[n].Comma.In(c.ctx)
+}
+
+func (c commas[T, _]) AppendComma(value T, comma token.Token) {
+	c.InsertComma(c.Len(), value, comma)
+}
+
+func (c commas[T, _]) InsertComma(n int, value T, comma token.Token) {
+	c.ctx.Nodes().panicIfNotOurs(comma)
+	v := c.SliceInserter.Unwrap(value)
+	v.Comma = comma.ID()
+
+	*c.Slice = slices.Insert(*c.Slice, n, v)
+}

experimental/ast/decl_body.go

@@ -15,9 +15,8 @@
 package ast
 
 import (
-	"slices"
-
 	"github.com/bufbuild/protocompile/experimental/report"
+	"github.com/bufbuild/protocompile/experimental/seq"
 	"github.com/bufbuild/protocompile/experimental/token"
 	"github.com/bufbuild/protocompile/internal/arena"
 )
@@ -28,8 +27,6 @@ import (
 // "orphaned" field or oneof outside of a message, or an RPC method inside of an enum, and
 // so on.
 //
-// DeclBody implements [Slice], providing access to its declarations.
-//
 // # Grammar
 //
 //	DeclBody := `{` DeclAny* `}`
@@ -48,10 +45,6 @@ type rawDeclBody struct {
 	ptrs  []arena.Untyped
 }
 
-var (
-	_ Inserter[DeclAny] = DeclBody{}
-)
-
 // Braces returns this body's surrounding braces, if it has any.
 func (d DeclBody) Braces() token.Token {
 	if d.IsZero() {
@@ -63,64 +56,39 @@ func (d DeclBody) Braces() token.Token {
 
 // Span implements [report.Spanner].
 func (d DeclBody) Span() report.Span {
+	decls := d.Decls()
 	switch {
 	case d.IsZero():
 		return report.Span{}
 	case !d.Braces().IsZero():
 		return d.Braces().Span()
-	case d.Len() == 0:
+	case decls.Len() == 0:
 		return report.Span{}
 	default:
-		return report.Join(d.At(0), d.At(d.Len()-1))
+		return report.Join(decls.At(0), decls.At(decls.Len()-1))
 	}
 }
 
-// Len returns the number of declarations inside of this body.
-func (d DeclBody) Len() int {
-	if d.IsZero() {
-		return 0
-	}
-
-	return len(d.raw.ptrs)
-}
-
-// At returns the nth element of this body.
-func (d DeclBody) At(n int) DeclAny {
-	return rawDecl{d.raw.ptrs[n], d.raw.kinds[n]}.With(d.Context())
-}
-
-// Iter is an iterator over the nodes inside this body.
-func (d DeclBody) Iter(yield func(int, DeclAny) bool) {
+// Decls returns a [seq.Inserter] over the declarations in this body.
+func (d DeclBody) Decls() seq.Inserter[DeclAny] {
+	type slice = seq.SliceInserter2[DeclAny, DeclKind, arena.Untyped]
 	if d.IsZero() {
-		return
+		return slice{}
 	}
 
-	for i := range d.raw.kinds {
-		if !yield(i, d.At(i)) {
-			break
-		}
+	return seq.SliceInserter2[DeclAny, DeclKind, arena.Untyped]{
+		Slice1: &d.raw.kinds,
+		Slice2: &d.raw.ptrs,
+		Wrap: func(k DeclKind, p arena.Untyped) DeclAny {
+			return rawDecl{p, k}.With(d.Context())
+		},
+		Unwrap: func(d DeclAny) (DeclKind, arena.Untyped) {
+			d.Context().Nodes().panicIfNotOurs(d)
+			return d.raw.kind, d.raw.ptr
+		},
 	}
 }
 
-// Append appends a new declaration to this body.
-func (d DeclBody) Append(value DeclAny) {
-	d.Insert(d.Len(), value)
-}
-
-// Insert inserts a new declaration at the given index.
-func (d DeclBody) Insert(n int, value DeclAny) {
-	d.Context().Nodes().panicIfNotOurs(value)
-
-	d.raw.kinds = slices.Insert(d.raw.kinds, n, value.Kind())
-	d.raw.ptrs = slices.Insert(d.raw.ptrs, n, value.raw.ptr)
-}
-
-// Delete deletes the declaration at the given index.
-func (d DeclBody) Delete(n int) {
-	d.raw.kinds = slices.Delete(d.raw.kinds, n, n+1)
-	d.raw.ptrs = slices.Delete(d.raw.ptrs, n, n+1)
-}
-
 func wrapDeclBody(c Context, ptr arena.Pointer[rawDeclBody]) DeclBody {
 	return DeclBody{wrapDecl(c, ptr)}
 }

experimental/ast/decl_file.go

@@ -16,6 +16,7 @@ package ast
 
 import (
 	"github.com/bufbuild/protocompile/experimental/report"
+	"github.com/bufbuild/protocompile/experimental/seq"
 	"github.com/bufbuild/protocompile/experimental/token"
 	"github.com/bufbuild/protocompile/internal/arena"
 	"github.com/bufbuild/protocompile/internal/iter"
@@ -36,7 +37,7 @@ type File struct {
 
 // Syntax returns this file's pragma, if it has one.
 func (f File) Syntax() (syntax DeclSyntax) {
-	f.Iter(func(_ int, d DeclAny) bool {
+	seq.Values(f.Decls())(func(d DeclAny) bool {
 		if s := d.AsSyntax(); !s.IsZero() {
 			syntax = s
 			return false
@@ -48,7 +49,7 @@ func (f File) Syntax() (syntax DeclSyntax) {
 
 // Package returns this file's package declaration, if it has one.
 func (f File) Package() (pkg DeclPackage) {
-	f.Iter(func(_ int, d DeclAny) bool {
+	seq.Values(f.Decls())(func(d DeclAny) bool {
 		if p := d.AsPackage(); !p.IsZero() {
 			pkg = p
 			return false
@@ -62,7 +63,7 @@ func (f File) Package() (pkg DeclPackage) {
 func (f File) Imports() iter.Seq2[int, DeclImport] {
 	return func(yield func(int, DeclImport) bool) {
 		var i int
-		f.Iter(func(_ int, d DeclAny) bool {
+		seq.Values(f.Decls())(func(d DeclAny) bool {
 			if imp := d.AsImport(); !imp.IsZero() {
 				if !yield(i, imp) {
 					return false

experimental/ast/decl_range.go

@@ -15,19 +15,15 @@
 package ast
 
 import (
-	"slices"
-
 	"github.com/bufbuild/protocompile/experimental/report"
+	"github.com/bufbuild/protocompile/experimental/seq"
 	"github.com/bufbuild/protocompile/experimental/token"
 	"github.com/bufbuild/protocompile/internal/arena"
 )
 
 // DeclRange represents an extension or reserved range declaration. They are almost identical
 // syntactically so they use the same AST node.
 //
-// In the Protocompile AST, ranges can contain arbitrary expressions. Thus, DeclRange
-// implements [Comma[ExprAny]].
-//
 // # Grammar
 //
 //	DeclRange := (`extensions` | `reserved`) (Expr `,`)* Expr? CompactOptions? `;`?
@@ -47,10 +43,6 @@ type DeclRangeArgs struct {
 	Semicolon token.Token
 }
 
-var (
-	_ Commas[ExprAny] = DeclRange{}
-)
-
 // Keyword returns the keyword for this range.
 func (d DeclRange) Keyword() token.Token {
 	if d.IsZero() {
@@ -70,64 +62,28 @@ func (d DeclRange) IsReserved() bool {
 	return d.Keyword().Text() == "reserved"
 }
 
-// Len implements [Slice].
-func (d DeclRange) Len() int {
-	if d.IsZero() {
-		return 0
-	}
-
-	return len(d.raw.args)
-}
-
-// At implements [Slice].
-func (d DeclRange) At(n int) ExprAny {
-	return newExprAny(d.Context(), d.raw.args[n].Value)
-}
-
-// Iter implements [Slice].
-func (d DeclRange) Iter(yield func(int, ExprAny) bool) {
+// Ranges returns the sequence of expressions denoting the ranges in this
+// range declaration.
+func (d DeclRange) Ranges() Commas[ExprAny] {
+	type slice = commas[ExprAny, rawExpr]
 	if d.IsZero() {
-		return
+		return slice{}
 	}
-	for i, arg := range d.raw.args {
-		if !yield(i, newExprAny(d.Context(), arg.Value)) {
-			break
-		}
+	return slice{
+		ctx: d.Context(),
+		SliceInserter: seq.SliceInserter[ExprAny, withComma[rawExpr]]{
+			Slice: &d.raw.args,
+			Wrap: func(c withComma[rawExpr]) ExprAny {
+				return newExprAny(d.Context(), c.Value)
+			},
+			Unwrap: func(e ExprAny) withComma[rawExpr] {
+				d.Context().Nodes().panicIfNotOurs(e)
+				return withComma[rawExpr]{Value: e.raw}
+			},
+		},
 	}
 }
 
-// Append implements [Inserter].
-func (d DeclRange) Append(expr ExprAny) {
-	d.InsertComma(d.Len(), expr, token.Zero)
-}
-
-// Insert implements [Inserter].
-func (d DeclRange) Insert(n int, expr ExprAny) {
-	d.InsertComma(n, expr, token.Zero)
-}
-
-// Delete implements [Inserter].
-func (d DeclRange) Delete(n int) {
-	d.raw.args = slices.Delete(d.raw.args, n, n+1)
-}
-
-// Comma implements [Commas].
-func (d DeclRange) Comma(n int) token.Token {
-	return d.raw.args[n].Comma.In(d.Context())
-}
-
-// AppendComma implements [Commas].
-func (d DeclRange) AppendComma(expr ExprAny, comma token.Token) {
-	d.InsertComma(d.Len(), expr, comma)
-}
-
-// InsertComma implements [Commas].
-func (d DeclRange) InsertComma(n int, expr ExprAny, comma token.Token) {
-	d.Context().Nodes().panicIfNotOurs(expr, comma)
-
-	d.raw.args = slices.Insert(d.raw.args, n, withComma[rawExpr]{expr.raw, comma.ID()})
-}
-
 // Options returns the compact options list for this range.
 func (d DeclRange) Options() CompactOptions {
 	if d.IsZero() {
@@ -157,16 +113,17 @@ func (d DeclRange) Semicolon() token.Token {
 
 // Span implements [report.Spanner].
 func (d DeclRange) Span() report.Span {
+	r := d.Ranges()
 	switch {
 	case d.IsZero():
 		return report.Span{}
-	case d.Len() == 0:
+	case r.Len() == 0:
 		return report.Join(d.Keyword(), d.Semicolon(), d.Options())
 	default:
 		return report.Join(
 			d.Keyword(), d.Semicolon(), d.Options(),
-			d.At(0),
-			d.At(d.Len()-1),
+			r.At(0),
+			r.At(r.Len()-1),
 		)
 	}
 }