/
rewrite.go
95 lines (85 loc) · 1.93 KB
/
rewrite.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
package rewrite
import (
"bytes"
"fmt"
"go/token"
)
func New(data []byte, base int) *Rewriter {
return &Rewriter{
segs: []segment{{
data: data,
start: base,
end: base + len(data),
}},
}
}
type Rewriter struct {
segs []segment
}
func (r *Rewriter) Replace(start, end token.Pos, data []byte) {
si, so := r.seg(start)
ei, eo := r.seg(end)
r.replace(si, so, ei, eo, data)
}
func (r *Rewriter) Insert(start token.Pos, data []byte) {
si, so := r.seg(start)
r.replace(si, so, si, so, data)
}
func (r *Rewriter) Delete(start, end token.Pos) {
si, so := r.seg(start)
ei, eo := r.seg(end)
r.replace(si, so, ei, eo, nil)
}
func (r *Rewriter) Data() []byte {
var buf bytes.Buffer
for _, seg := range r.segs {
buf.Write(seg.data)
}
return buf.Bytes()
}
func (r *Rewriter) replace(si, so, ei, eo int, data []byte) {
if si == ei {
// Same segment; cut it into two
start := r.segs[si]
end := segment{
start: start.start + eo,
end: start.end,
data: start.data[eo:],
}
start.data = start.data[:so]
start.end = start.start + so
mid := segment{
start: start.end,
end: end.start,
data: data,
}
r.segs = append(r.segs[:si], append([]segment{start, mid, end}, r.segs[ei+1:]...)...)
} else {
// Already different segments; adjust start/end and replace segments in-between
start := r.segs[si]
end := r.segs[ei]
start.end = start.start + so
start.data = start.data[:so]
end.start += eo
end.data = end.data[eo:]
mid := segment{
start: start.end,
end: end.start,
data: data,
}
r.segs = append(r.segs[:si], append([]segment{start, mid, end}, r.segs[ei+1:]...)...)
}
}
func (r *Rewriter) seg(pos token.Pos) (idx int, offset int) {
p := int(pos)
for i, seg := range r.segs {
if seg.start <= p && p < seg.end {
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
}