diff --git a/go/vt/sqlparser/pathbuilder/ast_paths_builder.go b/go/vt/sqlparser/pathbuilder/ast_paths_builder.go new file mode 100644 index 00000000000..9996ee8e722 --- /dev/null +++ b/go/vt/sqlparser/pathbuilder/ast_paths_builder.go @@ -0,0 +1,81 @@ +/* +Copyright 2025 The Vitess Authors. + +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 pathbuilder + +import "encoding/binary" + +// ASTPathBuilder is used to build +// paths for an AST. The steps are uints. +type ASTPathBuilder struct { + path []byte + sizes []int +} + +// NewASTPathBuilder creates a new ASTPathBuilder. +func NewASTPathBuilder() *ASTPathBuilder { + return &ASTPathBuilder{} +} + +// ToPath converts the ASTPathBuilder to a string. +func (apb *ASTPathBuilder) ToPath() string { + return string(apb.path) +} + +// AddStep appends a single step (2 bytes) to path. +func (apb *ASTPathBuilder) AddStep(step uint16) { + b := make([]byte, 2) + binary.BigEndian.PutUint16(b, step) + apb.path = append(apb.path, b...) + apb.sizes = append(apb.sizes, 2) +} + +// AddStepWithOffset appends a single step (2 bytes) to path +func (apb *ASTPathBuilder) AddStepWithOffset(step uint16, offset int) { + b := make([]byte, 2) + binary.BigEndian.PutUint16(b, step) + apb.path = append(apb.path, b...) + + b = make([]byte, 8) + bytesWritten := binary.PutVarint(b, int64(offset)) + apb.path = append(apb.path, b[:bytesWritten]...) + apb.sizes = append(apb.sizes, 2+bytesWritten) +} + +// ChangeOffset changes the offset of the last step +func (apb *ASTPathBuilder) ChangeOffset(newOffset int) { + previousOffsetSize := apb.sizes[len(apb.sizes)-1] - 2 + + b := make([]byte, 8) + bytesWritten := binary.PutVarint(b, int64(newOffset)) + + // Remove the previous offset + apb.path = apb.path[:len(apb.path)-previousOffsetSize] + // Add the new offset + apb.path = append(apb.path, b[:bytesWritten]...) + apb.sizes[len(apb.sizes)-1] = 2 + bytesWritten +} + +// Pop removes the last step from the path. +func (apb *ASTPathBuilder) Pop() { + if len(apb.sizes) == 0 { + return + } + + lastSize := apb.sizes[len(apb.sizes)-1] + apb.path = apb.path[:len(apb.path)-lastSize] + apb.sizes = apb.sizes[:len(apb.sizes)-1] +} diff --git a/go/vt/sqlparser/pathbuilder/ast_paths_builder_test.go b/go/vt/sqlparser/pathbuilder/ast_paths_builder_test.go new file mode 100644 index 00000000000..efea7a67fe3 --- /dev/null +++ b/go/vt/sqlparser/pathbuilder/ast_paths_builder_test.go @@ -0,0 +1,167 @@ +/* +Copyright 2025 The Vitess Authors. + +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 pathbuilder + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestASTPathBuilderAddStep(t *testing.T) { + tests := []struct { + name string + steps []uint16 + wantPath string + }{ + { + name: "single step", + steps: []uint16{1}, + wantPath: "\x00\x01", + }, + { + name: "multiple steps", + steps: []uint16{1, 0x24, 0x913}, + wantPath: "\x00\x01\x00\x24\x09\x13", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + apb := NewASTPathBuilder() + for _, step := range tt.steps { + apb.AddStep(step) + } + require.Equal(t, tt.wantPath, apb.ToPath()) + }) + } +} + +func TestASTPathBuilderAddStepOffset(t *testing.T) { + tests := []struct { + name string + steps []uint16 + offsets []int + wantPath string + }{ + { + name: "single step", + steps: []uint16{1}, + offsets: []int{0}, + wantPath: "\x00\x01\x00", + }, + { + name: "multiple steps", + steps: []uint16{1, 0x24, 0x913}, + offsets: []int{0, -1, 2}, + wantPath: "\x00\x01\x00\x00\x24\x09\x13\x04", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + apb := NewASTPathBuilder() + for idx, step := range tt.steps { + if tt.offsets[idx] == -1 { + apb.AddStep(step) + } else { + apb.AddStepWithOffset(step, tt.offsets[idx]) + } + } + require.Equal(t, tt.wantPath, apb.ToPath()) + }) + } +} + +func TestASTPathBuilderChangeOffset(t *testing.T) { + tests := []struct { + name string + steps []uint16 + offsets []int + changeOffsetValue int + wantPath string + }{ + { + name: "single step", + steps: []uint16{1}, + offsets: []int{0}, + changeOffsetValue: 5, + wantPath: "\x00\x01\x0a", + }, + { + name: "multiple steps", + steps: []uint16{1, 0x24, 0x913}, + offsets: []int{0, -1, 2}, + changeOffsetValue: 5, + wantPath: "\x00\x01\x00\x00\x24\x09\x13\x0a", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + apb := NewASTPathBuilder() + for idx, step := range tt.steps { + if tt.offsets[idx] == -1 { + apb.AddStep(step) + } else { + apb.AddStepWithOffset(step, tt.offsets[idx]) + } + } + apb.ChangeOffset(tt.changeOffsetValue) + require.Equal(t, tt.wantPath, apb.ToPath()) + }) + } +} + +func TestASTPathBuilderPop(t *testing.T) { + tests := []struct { + name string + steps []uint16 + offsets []int + wantPath string + }{ + { + name: "single step", + steps: []uint16{1}, + offsets: []int{0}, + wantPath: "", + }, + { + name: "multiple steps - final step with offset", + steps: []uint16{1, 0x24, 0x913}, + offsets: []int{0, -1, 2}, + wantPath: "\x00\x01\x00\x00\x24", + }, + { + name: "multiple steps - final step without offset", + steps: []uint16{1, 0x24, 0x913}, + offsets: []int{0, -1, -1}, + wantPath: "\x00\x01\x00\x00\x24", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + apb := NewASTPathBuilder() + for idx, step := range tt.steps { + if tt.offsets[idx] == -1 { + apb.AddStep(step) + } else { + apb.AddStepWithOffset(step, tt.offsets[idx]) + } + } + apb.Pop() + require.Equal(t, tt.wantPath, apb.ToPath()) + }) + } +}