diff --git a/v2/codegen/internal/rewrite/rewrite.go b/v2/codegen/internal/rewrite/rewrite.go index 0b88c1ace0..eb757c3af6 100644 --- a/v2/codegen/internal/rewrite/rewrite.go +++ b/v2/codegen/internal/rewrite/rewrite.go @@ -34,6 +34,17 @@ func (r *Rewriter) ReplaceNode(node ast.Node, data []byte) { } func (r *Rewriter) Insert(start token.Pos, data []byte) { + // If the pos is at the very end of the file, insert a new segment directly, + // since calling r.seg(start) would panic. + if len(r.segs) > 0 && r.segs[len(r.segs)-1].end == int(start) { + r.segs = append(r.segs, segment{ + start: int(start), + end: int(start) + len(data), + data: data, + }) + return + } + si, so := r.seg(start) r.replace(si, so, si, so, data) } @@ -93,10 +104,12 @@ func (r *Rewriter) seg(pos token.Pos) (idx int, offset int) { return i, int(p - seg.start) } } + panic(fmt.Sprintf("original file does not contain pos %v", pos)) } type segment struct { - start, end int - data []byte + start int // inclusive + end int // exclusive + data []byte } diff --git a/v2/codegen/internal/rewrite/rewrite_test.go b/v2/codegen/internal/rewrite/rewrite_test.go index c346914366..42938c799e 100644 --- a/v2/codegen/internal/rewrite/rewrite_test.go +++ b/v2/codegen/internal/rewrite/rewrite_test.go @@ -43,3 +43,15 @@ func TestDelete(t *testing.T) { t.Errorf("got data %s, want %s", got, want) } } + +func TestInsertAtEnd(t *testing.T) { + rw := New([]byte(""), 1) + rw.Insert(1, []byte("// test")) + if got, want := rw.Data(), []byte("// test"); !bytes.Equal(got, want) { + t.Errorf("got data %s, want %s", got, want) + } + rw.Delete(1, 4) + if got, want := rw.Data(), []byte("test"); !bytes.Equal(got, want) { + t.Errorf("got data %s, want %s", got, want) + } +}