Files
dolthub--doltgresql/testing/generation/utils/generator_stack.go
T
2026-07-13 12:32:25 +08:00

300 lines
11 KiB
Go

// Copyright 2023-2024 Dolthub, 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 utils
import (
"github.com/cockroachdb/errors"
"github.com/dolthub/doltgresql/utils"
)
// StatementGeneratorStackElement represents an element within the StatementGeneratorStack.
type StatementGeneratorStackElement struct {
depth int
generators []StatementGenerator
}
// StatementGeneratorStack handles the creation of a StatementGenerator by contextually applying operations based on the
// current state of the internal stack whenever calls are made.
type StatementGeneratorStack struct {
stack *utils.Stack[*StatementGeneratorStackElement]
depth int
}
// NewStatementGeneratorStack returns a new *StatementGeneratorStack.
func NewStatementGeneratorStack() *StatementGeneratorStack {
sgs := &StatementGeneratorStack{
stack: utils.NewStack[*StatementGeneratorStackElement](),
depth: 0,
}
sgs.stack.Push(NewStatementGeneratorStackElement(0))
return sgs
}
// NewStatementGeneratorStackElement returns a new *StatementGeneratorStackElement.
func NewStatementGeneratorStackElement(depth int, gens ...StatementGenerator) *StatementGeneratorStackElement {
return &StatementGeneratorStackElement{
depth: depth,
generators: gens,
}
}
// AddText creates a new TextGen at the current depth.
func (sgs *StatementGeneratorStack) AddText(text string) {
sgs.stack.Peek().Append(Text(text))
}
// AddVariable creates a new VariableGen at the current depth.
func (sgs *StatementGeneratorStack) AddVariable(name string) {
sgs.stack.Peek().Append(Variable(name, nil))
}
// Or will take all items from the current depth and add them to a parent OrGen. Either the previous depth is an OrGen,
// or the stack will increment the sub depth to insert an OrGen.
func (sgs *StatementGeneratorStack) Or() error {
if sgs.stack.Empty() {
return errors.Errorf("cannot apply Or to an empty stack")
}
parentElement := sgs.stack.PeekDepth(1)
current := sgs.aggregate(sgs.stack.Pop().generators)
if parentElement == nil {
// We're the root, so we put an OrGen as the root, and make the current a child of the new OrGen
sgs.stack.Push(NewStatementGeneratorStackElement(sgs.depth, Or(current)))
sgs.stack.Push(NewStatementGeneratorStackElement(sgs.depth))
} else if len(parentElement.generators) == 0 {
if parentElement.depth == sgs.depth {
// We're still at the same depth, so it's safe to append an OrGen
parentElement.Append(Or(current))
} else {
// We should retain the depth, so we need to create a new element since there's a depth boundary
sgs.stack.Push(NewStatementGeneratorStackElement(sgs.depth, Or(current)))
}
sgs.stack.Push(NewStatementGeneratorStackElement(sgs.depth))
} else {
if parentElement.depth == sgs.depth {
// We're still at the same depth, so we need to check if the parent is an OrGen or not
if orGen, ok := parentElement.LastGenerator().(*OrGen); ok {
// The parent is an OrGen, so we can simply add this as another option
if err := orGen.AddChildren(current); err != nil {
return err
}
} else {
// The parent is not an OrGen, so we need to add to the depth
sgs.stack.Push(NewStatementGeneratorStackElement(sgs.depth, Or(current)))
}
} else {
// We should retain the depth, so we need to create a new element since there's a depth boundary
sgs.stack.Push(NewStatementGeneratorStackElement(sgs.depth, Or(current)))
}
sgs.stack.Push(NewStatementGeneratorStackElement(sgs.depth))
}
return nil
}
// NewScope increases the depth.
func (sgs *StatementGeneratorStack) NewScope() {
sgs.depth++
sgs.stack.Push(NewStatementGeneratorStackElement(sgs.depth))
}
// NewOptionalScope increases the depth, creating an OptionalGen at the depth root, and adding elements to the sub depth.
func (sgs *StatementGeneratorStack) NewOptionalScope() {
sgs.depth++
sgs.stack.Push(NewStatementGeneratorStackElement(sgs.depth, Optional()))
sgs.stack.Push(NewStatementGeneratorStackElement(sgs.depth))
}
// NewParenScope increases the depth, creating a CollectionGen at the depth root, adding Text("(") to the CollectionGen,
// and adding any new elements to the sub depth.
func (sgs *StatementGeneratorStack) NewParenScope() {
sgs.depth++
sgs.stack.Push(NewStatementGeneratorStackElement(sgs.depth, Collection(Text("("))))
sgs.stack.Push(NewStatementGeneratorStackElement(sgs.depth))
}
// ExitScope decrements the depth, writing all elements from the current depth (and all sub depths) to the preceding
// depth.
func (sgs *StatementGeneratorStack) ExitScope() error {
defer func() {
sgs.depth--
}()
current := sgs.aggregate(sgs.stack.Pop().generators)
for parentElement := sgs.stack.Peek(); parentElement.depth == sgs.depth; parentElement = sgs.stack.Peek() {
if orGen, ok := parentElement.LastGenerator().(*OrGen); ok {
if err := orGen.AddChildren(current); err != nil {
return err
}
} else {
parentElement.Append(current)
}
current = sgs.aggregate(sgs.stack.Pop().generators)
}
sgs.stack.Peek().Append(current)
return nil
}
// ExitOptionalScope decrements the depth, writing all elements from the current depth (and all sub depths) to the
// preceding depth. This will fail if the root of the current depth is not an OptionalGen.
func (sgs *StatementGeneratorStack) ExitOptionalScope() error {
defer func() {
sgs.depth--
}()
current := sgs.aggregate(sgs.stack.Pop().generators)
for parentElement := sgs.stack.Peek(); sgs.stack.PeekDepth(1).depth == sgs.depth; parentElement = sgs.stack.Peek() {
if orGen, ok := parentElement.LastGenerator().(*OrGen); ok {
if err := orGen.AddChildren(current); err != nil {
return err
}
} else {
parentElement.Append(current)
}
current = sgs.aggregate(sgs.stack.Pop().generators)
}
optionalElement := sgs.stack.Pop()
if optionalElement.depth != sgs.depth {
return errors.Errorf("internal bookkeeping error, attempted to exit from optional scope but the depth is incorrect")
}
if optionalGen, ok := optionalElement.LastGenerator().(*OptionalGen); ok {
if err := optionalGen.AddChildren(current); err != nil {
return err
}
} else {
return errors.Errorf("internal bookkeeping error, attempted to exit from optional scope but missing OptionalGen")
}
sgs.stack.Peek().Append(sgs.aggregate(optionalElement.generators))
return nil
}
// ExitParenScope decrements the depth, writing all elements from the current depth (and all sub depths) to the
// preceding depth, while appending an ending Text(")").
func (sgs *StatementGeneratorStack) ExitParenScope() error {
defer func() {
sgs.depth--
}()
current := sgs.aggregate(sgs.stack.Pop().generators)
for parentElement := sgs.stack.Peek(); sgs.stack.PeekDepth(1).depth == sgs.depth; parentElement = sgs.stack.Peek() {
if orGen, ok := parentElement.LastGenerator().(*OrGen); ok {
if err := orGen.AddChildren(current); err != nil {
return err
}
} else {
parentElement.Append(current)
}
current = sgs.aggregate(sgs.stack.Pop().generators)
}
collectionElement := sgs.stack.Pop()
if collectionElement.depth != sgs.depth {
return errors.Errorf("internal bookkeeping error, attempted to exit from paren scope but the depth is incorrect")
}
if collectionGen, ok := collectionElement.LastGenerator().(*CollectionGen); ok {
if err := collectionGen.AddChildren(current, Text(")")); err != nil {
return err
}
} else {
return errors.Errorf("internal bookkeeping error, attempted to exit from paren scope but missing CollectionGen")
}
sgs.stack.Peek().Append(sgs.aggregate(collectionElement.generators))
return nil
}
// Repeat will add the last StatementGenerator at the current depth to an OptionalGen.
func (sgs *StatementGeneratorStack) Repeat() error {
current := sgs.stack.Peek()
lastGen := current.LastGenerator()
if lastGen == nil {
return errors.Errorf("unable to repeat as no generators exist at the current depth")
}
current.Append(Optional(lastGen.Copy()))
return nil
}
// OptionalRepeat will add the last StatementGenerator at the current depth to an OptionalGen. If a prefix is given,
// then it will be added as a TextGen before the repeated generator.
func (sgs *StatementGeneratorStack) OptionalRepeat(prefix string) error {
current := sgs.stack.Peek()
lastGen := current.LastGenerator()
if lastGen == nil {
return errors.Errorf("unable to optionally repeat as no generators exist at the current depth")
}
if len(prefix) > 0 {
current.Append(Optional(Collection(Text(prefix), lastGen.Copy())))
} else {
current.Append(Optional(lastGen.Copy()))
}
return nil
}
// Finish returns a StatementGenerator containing all of the generators that have been added. This returns an error if
// the stack is not at the root depth, as that indicates a missing scope exit, or too many scope entries.
func (sgs *StatementGeneratorStack) Finish() (StatementGenerator, error) {
if sgs.depth != 0 {
if sgs.depth < 0 {
return nil, errors.Errorf("depth is invalid, too many scope exits")
} else {
return nil, errors.Errorf("depth is invalid, too many scope entries")
}
} else if !sgs.stack.Empty() && sgs.stack.Peek().depth != 0 {
return nil, errors.Errorf("internal bookkeeping error, stack depth does not match handle depth")
}
if sgs.stack.Len() == 1 && len(sgs.stack.Peek().generators) == 0 {
return nil, nil
}
var lastDepth []StatementGenerator
for !sgs.stack.Empty() {
currentDepth := sgs.stack.Pop()
if len(currentDepth.generators) == 0 {
return nil, errors.Errorf("internal bookkeeping error, stack has a depth with no generators")
}
if lastGen := currentDepth.LastGenerator(); lastGen != nil {
if orGen, ok := lastGen.(*OrGen); ok {
if err := orGen.AddChildren(sgs.aggregate(lastDepth)); err != nil {
return nil, err
}
} else {
currentDepth.Append(sgs.aggregate(lastDepth))
}
} else {
currentDepth.Append(sgs.aggregate(lastDepth))
}
lastDepth = currentDepth.generators
}
return sgs.aggregate(lastDepth), nil
}
// aggregate returns an aggregate of the given generators. Returns nil if the slice is empty.
func (sgs *StatementGeneratorStack) aggregate(gens []StatementGenerator) StatementGenerator {
gens = removeNilGenerators(gens)
if len(gens) == 0 {
return nil
} else if len(gens) == 1 {
return gens[0]
} else {
return Collection(gens...)
}
}
// LastGenerator returns the last StatementGenerator contained within its slice.
func (sgse *StatementGeneratorStackElement) LastGenerator() StatementGenerator {
if sgse == nil || len(sgse.generators) == 0 {
return nil
}
return sgse.generators[len(sgse.generators)-1]
}
// Append adds the given child to this element's collection.
func (sgse *StatementGeneratorStackElement) Append(child StatementGenerator) {
sgse.generators = append(sgse.generators, child)
}