Skip to content

Commit a429cb3

Browse files
AsterDYAsterDY
committed
feat: support Patcher
1 parent 0fb56ca commit a429cb3

File tree

7 files changed

+291
-8
lines changed

7 files changed

+291
-8
lines changed

‎src/compress/golang/plugin/parse/parser.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -147,7 +147,7 @@ func (p *GoParser) ParseModule(mod *Module, dir string) (err error) {
147147
return nil
148148
}
149149
rel, _ := filepath.Rel(p.homePageDir, path)
150-
mod.Files[rel] = NewFile(path)
150+
mod.Files[rel] = NewFile(rel)
151151
return nil
152152
})
153153
return p.loadPackages(mod, dir, "./...")

‎src/lang/collect/export.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -141,7 +141,7 @@ func (c *Collector) exportSymbol(repo *uniast.Repository, symbol *DocumentSymbol
141141
fileLine := c.fileLine(symbol.Location)
142142
// collect files
143143
if module.Files[relfile] == nil {
144-
module.Files[relfile] = uniast.NewFile(file)
144+
module.Files[relfile] = uniast.NewFile(relfile)
145145
}
146146

147147
content := symbol.Text

‎src/lang/patch/lib.go

Lines changed: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,133 @@
1+
// Copyright 2025 ByteDance Inc.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// https://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package patch
16+
17+
import (
18+
"fmt"
19+
"math"
20+
"os"
21+
"path/filepath"
22+
"sort"
23+
24+
"github.com/cloudwego/abcoder/src/lang/utils"
25+
"github.com/cloudwego/abcoder/src/uniast"
26+
)
27+
28+
// PatchModule patches the ast Nodes onto module files
29+
30+
type Patch struct {
31+
Id uniast.Identity
32+
Codes string
33+
File string
34+
Type uniast.NodeType
35+
}
36+
37+
type patchNode struct {
38+
uniast.FileLine
39+
Codes string
40+
}
41+
42+
type Patcher struct {
43+
Options
44+
repo *uniast.Repository
45+
patches map[string][]patchNode
46+
}
47+
48+
type Options struct {
49+
RepoDir string
50+
OutDir string
51+
}
52+
53+
func NewPatcher(repo *uniast.Repository, opts Options) *Patcher {
54+
return &Patcher{
55+
Options: opts,
56+
repo: repo,
57+
}
58+
}
59+
60+
func (p *Patcher) Patch(patch Patch) error {
61+
// find package
62+
node := p.repo.GetNode(patch.Id)
63+
if node == nil {
64+
node = p.repo.SetNode(patch.Id, patch.Type)
65+
}
66+
f := p.repo.GetFile(patch.Id)
67+
if f == nil {
68+
mod := p.repo.GetModule(patch.Id.ModPath)
69+
f = uniast.NewFile(patch.File)
70+
mod.SetFile(patch.File, f)
71+
node.SetFile(patch.File)
72+
}
73+
if err := p.patchFile(patchNode{FileLine: node.FileLine(), Codes: patch.Codes}); err != nil {
74+
return fmt.Errorf("patch file %s failed: %v", f.Path, err)
75+
}
76+
return nil
77+
}
78+
79+
func (p *Patcher) patchFile(n patchNode) error {
80+
if p.patches == nil {
81+
p.patches = make(map[string][]patchNode)
82+
}
83+
if n.StartOffset < 1 {
84+
n.StartOffset = math.MaxInt
85+
}
86+
p.patches[n.FileLine.File] = append(p.patches[n.FileLine.File], n)
87+
return nil
88+
}
89+
90+
func (p *Patcher) Flush() error {
91+
// write pathes
92+
for fpath, ns := range p.patches {
93+
sort.SliceStable(ns, func(i, j int) bool {
94+
return ns[i].StartOffset < ns[j].StartOffset
95+
})
96+
path := filepath.Join(p.RepoDir, fpath)
97+
data, err := os.ReadFile(path)
98+
if err != nil {
99+
return fmt.Errorf("read file %s failed: %v", path, err)
100+
}
101+
var offset int
102+
for _, n := range ns {
103+
if n.StartOffset >= len(data) {
104+
data = append(append(data, '\n'), []byte(n.Codes)...)
105+
continue
106+
}
107+
tmp := append(data[:offset+n.StartOffset:offset+n.StartOffset], []byte(n.Codes)...)
108+
data = append(tmp, data[offset+n.EndOffset:]...)
109+
offset += (len(n.Codes) - (n.EndOffset - n.StartOffset))
110+
}
111+
if err := utils.MustWriteFile(filepath.Join(p.OutDir, fpath), data); err != nil {
112+
return fmt.Errorf("write file %s failed: %v", fpath, err)
113+
}
114+
}
115+
// write origins
116+
for _, mod := range p.repo.Modules {
117+
for _, f := range mod.Files {
118+
if p.patches[f.Path] != nil {
119+
continue
120+
}
121+
fpath := filepath.Join(p.RepoDir, f.Path)
122+
bs, err := os.ReadFile(fpath)
123+
if err != nil {
124+
return fmt.Errorf("read file %s failed: %v", fpath, err)
125+
}
126+
fpath = filepath.Join(p.OutDir, f.Path)
127+
if err := utils.MustWriteFile(fpath, bs); err != nil {
128+
return fmt.Errorf("write file %s failed: %v", fpath, err)
129+
}
130+
}
131+
}
132+
return nil
133+
}

‎src/lang/patch/lib_test.go

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,79 @@
1+
/**
2+
* Copyright 2025 ByteDance Inc.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package patch
18+
19+
import (
20+
"testing"
21+
22+
"github.com/cloudwego/abcoder/src/uniast"
23+
)
24+
25+
var root = "../../../tmp"
26+
27+
func TestPatcher(t *testing.T) {
28+
// Load repository
29+
repo, err := uniast.LoadRepo(root + "/localsession.json")
30+
if err != nil {
31+
t.Errorf("failed to load repo: %v", err)
32+
}
33+
34+
// Create patcher with options
35+
patcher := NewPatcher(repo, Options{
36+
RepoDir: root + "/localsession",
37+
OutDir: root + "/localsession2",
38+
})
39+
40+
// Create a test patch
41+
testPatches := []Patch{
42+
{
43+
Id: uniast.Identity{ModPath: "github.com/cloudwego/localsession", PkgPath: "github.com/cloudwego/localsession/backup", Name: "DefaultOptions"},
44+
Codes: `func DefaultOptions() Options {
45+
ret := Options{
46+
Enable: false,
47+
ManagerOptions: localsession.DefaultManagerOptions(),
48+
}
49+
return ret
50+
}`,
51+
File: "backup/metainfo.go",
52+
Type: uniast.FUNC,
53+
},
54+
{
55+
Id: uniast.Identity{ModPath: "github.com/cloudwego/localsession", PkgPath: "github.com/cloudwego/localsession/backup", Name: "DefaultOptions2"},
56+
Codes: `func DefaultOptions2() Options {
57+
ret := Options{
58+
Enable: false,
59+
ManagerOptions: localsession.DefaultManagerOptions(),
60+
}
61+
return ret
62+
}`,
63+
File: "backup/metainfo.go",
64+
Type: uniast.FUNC,
65+
},
66+
}
67+
68+
// Apply the patches
69+
for _, testPatch := range testPatches {
70+
if err := patcher.Patch(testPatch); err != nil {
71+
t.Errorf("failed to patch: %v", err)
72+
}
73+
}
74+
75+
// Flush changes
76+
if err := patcher.Flush(); err != nil {
77+
t.Errorf("failed to flush: %v", err)
78+
}
79+
}

‎src/lang/utils/files.go

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,11 @@
11
// Copyright 2025 CloudWeGo Authors
2-
//
2+
//
33
// Licensed under the Apache License, Version 2.0 (the "License");
44
// you may not use this file except in compliance with the License.
55
// You may obtain a copy of the License at
6-
//
6+
//
77
// https://www.apache.org/licenses/LICENSE-2.0
8-
//
8+
//
99
// Unless required by applicable law or agreed to in writing, software
1010
// distributed under the License is distributed on an "AS IS" BASIS,
1111
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -15,6 +15,7 @@
1515
package utils
1616

1717
import (
18+
"fmt"
1819
"os"
1920
"path/filepath"
2021
"strings"
@@ -53,3 +54,14 @@ func FirstFile(dir string, subfix string, skipdir string) string {
5354
})
5455
return ret
5556
}
57+
58+
func MustWriteFile(fpath string, data []byte) error {
59+
dir := filepath.Dir(fpath)
60+
if err := os.MkdirAll(dir, 0755); err != nil {
61+
return fmt.Errorf("mkdir %s failed: %v", dir, err)
62+
}
63+
if err := os.WriteFile(fpath, data, 0644); err != nil {
64+
return fmt.Errorf("write file %s failed: %v", fpath, err)
65+
}
66+
return nil
67+
}

‎src/uniast/ast.go

Lines changed: 34 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -53,14 +53,23 @@ type File struct {
5353
}
5454

5555
func NewFile(path string) *File {
56-
abs, _ := filepath.Abs(path)
56+
// abs, _ := filepath.Abs(path)
5757
ret := File{
5858
Name: filepath.Base(path),
59-
Path: abs,
59+
Path: path,
6060
}
6161
return &ret
6262
}
6363

64+
func (m Module) SetFile(path string, file *File) {
65+
if m.Files == nil {
66+
m.Files = map[string]*File{}
67+
}
68+
if m.Files[path] == nil {
69+
m.Files[path] = file
70+
}
71+
}
72+
6473
type Module struct {
6574
Language Language
6675
Name string // go module name
@@ -70,6 +79,18 @@ type Module struct {
7079
Files map[string]*File `json:",omitempty"` // relative path => file info
7180
}
7281

82+
func (r Repository) GetFile(id Identity) *File {
83+
mod := r.Modules[id.ModPath]
84+
if mod == nil {
85+
return nil
86+
}
87+
node := r.GetNode(id)
88+
if node == nil {
89+
return nil
90+
}
91+
return mod.Files[node.FileLine().File]
92+
}
93+
7394
func IsExternalModule(modpath string) bool {
7495
return modpath == "" || strings.Contains(modpath, "@")
7596
}
@@ -112,6 +133,7 @@ type Package struct {
112133
Types map[string]*Type // type name => type define
113134
Vars map[string]*Var // var name => var define
114135
CompressData *string `json:"compress_data,omitempty"` // package compress info
136+
Path string // relative path to repo
115137
}
116138

117139
func NewPackage(pkgPath PkgPath) *Package {
@@ -127,9 +149,18 @@ func NewPackage(pkgPath PkgPath) *Package {
127149
// PkgPath is the import path of a package, it is either absolute path or url
128150
type PkgPath = string
129151

152+
type ModPath = string
153+
154+
func ModPathName(mod ModPath) string {
155+
if strings.Contains(mod, "@") {
156+
return strings.Split(mod, "@")[0]
157+
}
158+
return mod
159+
}
160+
130161
// Identity holds identity information about a third party declaration
131162
type Identity struct {
132-
ModPath string // ModPath is the module which the package belongs to
163+
ModPath // ModPath is the module which the package belongs to
133164
PkgPath // Import Path of the third party package
134165
Name string // Unique Name of declaration (FunctionName, TypeName.MethodName, InterfaceName<TypeName>.MethodName, or TypeName)
135166
}

‎src/uniast/node.go

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,26 @@ func (r *Repository) GetNode(id Identity) *Node {
2828
return node
2929
}
3030

31+
func (r *Repository) GetPackage(id Identity) *Package {
32+
mod, ok := r.Modules[id.ModPath]
33+
if !ok {
34+
return nil
35+
}
36+
pkg, ok := mod.Packages[id.PkgPath]
37+
if !ok {
38+
return nil
39+
}
40+
return pkg
41+
}
42+
43+
func (r *Repository) GetModule(mod ModPath) *Module {
44+
m, ok := r.Modules[mod]
45+
if !ok {
46+
return nil
47+
}
48+
return m
49+
}
50+
3151
// NOTICE: if entity not exist, only set the node on graph
3252
func (r *Repository) SetNode(id Identity, typ NodeType) *Node {
3353
key := id.Full()
@@ -215,6 +235,14 @@ type Node struct {
215235
Repo *Repository `json:"-"`
216236
}
217237

238+
func NewNode(id Identity, typ NodeType, repo *Repository) *Node {
239+
return &Node{
240+
Identity: id,
241+
Type: typ,
242+
Repo: repo,
243+
}
244+
}
245+
218246
type NodeID struct {
219247
Module string `json:"module,omitempty" jsonschema:"description=the building moudle of the ast node, e.g. github.com/bytedance/[email protected]"`
220248
Package string `json:"package" jsonschema:"description=the namespace of the ast node, e.g. github.com/bytedance/sonic/ast"`

0 commit comments

Comments
 (0)