blueprint/bpmodify/bpmodify.go

// Copyright 2024 Google Inc. All rights reserved.
//
// 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 bpmodify

import (
	"bytes"
	"errors"
	"fmt"
	"slices"
	"strings"

	"github.com/google/blueprint/parser"
)

// NewBlueprint returns a Blueprint for the given file contents that allows making modifications.
func NewBlueprint(filename string, data []byte) (*Blueprint, error) {
	r := bytes.NewReader(data)
	file, errs := parser.Parse(filename, r)
	if len(errs) > 0 {
		return nil, errors.Join(errs...)
	}

	return &Blueprint{
		data:   data,
		bpFile: file,
	}, nil
}

type Blueprint struct {
	data     []byte
	bpFile   *parser.File
	modified bool
}

// Bytes returns a copy of the current, possibly modified contents of the Blueprint as a byte slice.
func (bp *Blueprint) Bytes() ([]byte, error) {
	if bp.modified {
		data, err := parser.Print(bp.bpFile)
		if err != nil {
			return nil, err
		}
		return data, nil
	}
	return slices.Clone(bp.data), nil
}

// String returns the current, possibly modified contents of the Blueprint as a string.
func (bp *Blueprint) String() string {
	data, err := bp.Bytes()
	if err != nil {
		return err.Error()
	}
	return string(data)
}

// Modified returns true if any of the calls on the Blueprint caused the contents to be modified.
func (bp *Blueprint) Modified() bool {
	return bp.modified
}

// ModulesByName returns a ModuleSet that contains all modules with the given list of names.
// Requesting a module that does not exist is not an error.
func (bp *Blueprint) ModulesByName(names ...string) *ModuleSet {
	moduleSet := &ModuleSet{
		bp: bp,
	}
	for _, def := range bp.bpFile.Defs {
		module, ok := def.(*parser.Module)
		if !ok {
			continue
		}

		for _, prop := range module.Properties {
			if prop.Name == "name" {
				if stringValue, ok := prop.Value.(*parser.String); ok && slices.Contains(names, stringValue.Value) {
					moduleSet.modules = append(moduleSet.modules, module)
				}
			}
		}
	}

	return moduleSet
}

// AllModules returns a ModuleSet that contains all modules in the Blueprint.
func (bp *Blueprint) AllModules() *ModuleSet {
	moduleSet := &ModuleSet{
		bp: bp,
	}
	for _, def := range bp.bpFile.Defs {
		module, ok := def.(*parser.Module)
		if !ok {
			continue
		}

		moduleSet.modules = append(moduleSet.modules, module)
	}

	return moduleSet
}

// A ModuleSet represents a set of modules in a Blueprint, and can be used to make modifications
// the modules.
type ModuleSet struct {
	bp      *Blueprint
	modules []*parser.Module
}

// GetProperty returns a PropertySet that contains all properties with the given list of names
// in all modules in the ModuleSet.  Requesting properties that do not exist is not an error.
// It returns an error for a malformed property name, or if the requested property is nested
// in a property that is not a map.
func (ms *ModuleSet) GetProperty(properties ...string) (*PropertySet, error) {
	propertySet := &PropertySet{
		bp: ms.bp,
	}

	targetProperties, err := parseQualifiedProperties(properties)
	if err != nil {
		return nil, err
	}

	for _, targetProperty := range targetProperties {
		for _, module := range ms.modules {
			prop, _, err := getRecursiveProperty(module, targetProperty)
			if err != nil {
				return nil, err
			} else if prop == nil {
				continue
			}
			propertySet.properties = append(propertySet.properties, &property{
				property: prop,
				module:   module,
				name:     targetProperty,
			})
		}
	}

	return propertySet, nil
}

// GetOrCreateProperty returns a PropertySet that contains all properties with the given list of names
// in all modules in the ModuleSet, creating empty placeholder properties if they don't exist.
// It returns an error for a malformed property name, or if the requested property is nested
// in a property that is not a map.
func (ms *ModuleSet) GetOrCreateProperty(typ Type, properties ...string) (*PropertySet, error) {
	propertySet := &PropertySet{
		bp: ms.bp,
	}

	targetProperties, err := parseQualifiedProperties(properties)
	if err != nil {
		return nil, err
	}

	for _, targetProperty := range targetProperties {
		for _, module := range ms.modules {
			prop, _, err := getRecursiveProperty(module, targetProperty)
			if err != nil {
				return nil, err
			} else if prop == nil {
				prop, err = createRecursiveProperty(module, targetProperty, parser.ZeroExpression(parser.Type(typ)))
				if err != nil {
					return nil, err
				}
				ms.bp.modified = true
			} else {
				if prop.Value.Type() != parser.Type(typ) {
					return nil, fmt.Errorf("unexpected type found in property %q, wanted %s, found %s",
						targetProperty.String(), typ, prop.Value.Type())
				}
			}
			propertySet.properties = append(propertySet.properties, &property{
				property: prop,
				module:   module,
				name:     targetProperty,
			})
		}
	}

	return propertySet, nil
}

// RemoveProperty removes the given list of properties from all modules in the ModuleSet.
// It returns an error for a malformed property name, or if the requested property is nested
// in a property that is not a map.  Removing a property that does not exist is not an error.
func (ms *ModuleSet) RemoveProperty(properties ...string) error {
	targetProperties, err := parseQualifiedProperties(properties)
	if err != nil {
		return err
	}

	for _, targetProperty := range targetProperties {
		for _, module := range ms.modules {
			prop, parent, err := getRecursiveProperty(module, targetProperty)
			if err != nil {
				return err
			} else if prop != nil {
				parent.RemoveProperty(prop.Name)
				ms.bp.modified = true
			}
		}
	}
	return nil
}

// MoveProperty moves the given list of properties to a new parent property.
// It returns an error for a malformed property name, or if the requested property is nested
// in a property that is not a map.  Moving a property that does not exist is not an error.
func (ms *ModuleSet) MoveProperty(newParent string, properties ...string) error {
	targetProperties, err := parseQualifiedProperties(properties)
	if err != nil {
		return err
	}

	for _, targetProperty := range targetProperties {
		for _, module := range ms.modules {
			prop, parent, err := getRecursiveProperty(module, targetProperty)
			if err != nil {
				return err
			} else if prop != nil {
				parent.MovePropertyContents(prop.Name, newParent)
				ms.bp.modified = true
			}
		}
	}
	return nil
}

// PropertySet represents a set of properties in a set of modules.
type PropertySet struct {
	bp         *Blueprint
	properties []*property
	sortLists  bool
}

type property struct {
	property *parser.Property
	module   *parser.Module
	name     *qualifiedProperty
}

// SortListsWhenModifying causes any future modifications to lists in the PropertySet to sort
// the lists.  Otherwise, lists are only sorted if they appear to be sorted before modification.
func (ps *PropertySet) SortListsWhenModifying() {
	ps.sortLists = true
}

// SetString sets all properties in the PropertySet to the given string.  It returns an error
// if any of the properties are not strings.
func (ps *PropertySet) SetString(s string) error {
	var errs []error
	for _, prop := range ps.properties {
		value := prop.property.Value
		str, ok := value.(*parser.String)
		if !ok {
			errs = append(errs, fmt.Errorf("expected property %s in module %s to be string, found %s",
				prop.name, prop.module.Name(), value.Type().String()))
			continue
		}
		if str.Value != s {
			str.Value = s
			ps.bp.modified = true
		}
	}

	return errors.Join(errs...)
}

// SetBool sets all properties in the PropertySet to the given boolean.  It returns an error
// if any of the properties are not booleans.
func (ps *PropertySet) SetBool(b bool) error {
	var errs []error
	for _, prop := range ps.properties {
		value := prop.property.Value
		res, ok := value.(*parser.Bool)
		if !ok {
			errs = append(errs, fmt.Errorf("expected property %s in module %s to be bool, found %s",
				prop.name, prop.module.Name(), value.Type().String()))
			continue
		}
		if res.Value != b {
			res.Value = b
			ps.bp.modified = true
		}
	}
	return errors.Join(errs...)
}

// AddStringToList adds the given strings to all properties in the PropertySet.  It returns an error
// if any of the properties are not lists of strings.
func (ps *PropertySet) AddStringToList(strs ...string) error {
	var errs []error
	for _, prop := range ps.properties {
		value := prop.property.Value
		list, ok := value.(*parser.List)
		if !ok {
			errs = append(errs, fmt.Errorf("expected property %s in module %s to be list, found %s",
				prop.name, prop.module.Name(), value.Type()))
			continue
		}
		wasSorted := parser.ListIsSorted(list)
		modified := false
		for _, s := range strs {
			m := parser.AddStringToList(list, s)
			modified = modified || m
		}
		if modified {
			ps.bp.modified = true
			if wasSorted || ps.sortLists {
				parser.SortList(ps.bp.bpFile, list)
			}
		}
	}

	return errors.Join(errs...)
}

// RemoveStringFromList removes the given strings to all properties in the PropertySet if they are present.
// It returns an error  if any of the properties are not lists of strings.
func (ps *PropertySet) RemoveStringFromList(strs ...string) error {
	var errs []error
	for _, prop := range ps.properties {
		value := prop.property.Value
		list, ok := value.(*parser.List)
		if !ok {
			errs = append(errs, fmt.Errorf("expected property %s in module %s to be list, found %s",
				prop.name, prop.module.Name(), value.Type()))
			continue
		}
		wasSorted := parser.ListIsSorted(list)
		modified := false
		for _, s := range strs {
			m := parser.RemoveStringFromList(list, s)
			modified = modified || m
		}
		if modified {
			ps.bp.modified = true
			if wasSorted || ps.sortLists {
				parser.SortList(ps.bp.bpFile, list)
			}
		}
	}

	return errors.Join(errs...)
}

// AddLiteral adds the given literal blueprint snippet to all properties in the PropertySet if they are present.
// It returns an error  if any of the properties are not lists.
func (ps *PropertySet) AddLiteral(s string) error {
	var errs []error
	for _, prop := range ps.properties {
		value := prop.property.Value
		if ps.sortLists {
			return fmt.Errorf("sorting not supported when adding a literal")
		}
		list, ok := value.(*parser.List)
		if !ok {
			errs = append(errs, fmt.Errorf("expected property %s in module %s to be list, found %s",
				prop.name, prop.module.Name(), value.Type().String()))
			continue
		}
		value, parseErrs := parser.ParseExpression(strings.NewReader(s))
		if len(parseErrs) > 0 {
			errs = append(errs, parseErrs...)
			continue
		}
		list.Values = append(list.Values, value)
		ps.bp.modified = true
	}

	return errors.Join(errs...)
}

// ReplaceStrings applies replacements to all properties in the PropertySet.  It replaces all instances
// of the strings in the keys of the given map with their corresponding values.  It returns an error
// if any of the properties are not lists of strings.
func (ps *PropertySet) ReplaceStrings(replacements map[string]string) error {
	var errs []error
	for _, prop := range ps.properties {
		value := prop.property.Value
		if list, ok := value.(*parser.List); ok {
			modified := parser.ReplaceStringsInList(list, replacements)
			if modified {
				ps.bp.modified = true
			}
		} else if str, ok := value.(*parser.String); ok {
			oldVal := str.Value
			replacementValue := replacements[oldVal]
			if replacementValue != "" {
				str.Value = replacementValue
				ps.bp.modified = true
			}
		} else {
			errs = append(errs, fmt.Errorf("expected property %s in module %s to be a list or string, found %s",
				prop.name, prop.module.Name(), value.Type().String()))
		}
	}
	return errors.Join(errs...)
}

func getRecursiveProperty(module *parser.Module, property *qualifiedProperty) (prop *parser.Property,
	parent *parser.Map, err error) {

	parent, err = traverseToQualifiedPropertyParent(module, property, false)
	if err != nil {
		return nil, nil, err
	}
	if parent == nil {
		return nil, nil, nil
	}
	if prop, found := parent.GetProperty(property.name()); found {
		return prop, parent, nil
	}

	return nil, nil, nil
}

func createRecursiveProperty(module *parser.Module, property *qualifiedProperty,
	value parser.Expression) (prop *parser.Property, err error) {
	parent, err := traverseToQualifiedPropertyParent(module, property, true)
	if err != nil {
		return nil, err
	}
	if _, found := parent.GetProperty(property.name()); found {
		return nil, fmt.Errorf("property %q already exists", property.String())
	}

	prop = &parser.Property{Name: property.name(), Value: value}
	parent.Properties = append(parent.Properties, prop)
	return prop, nil
}

func traverseToQualifiedPropertyParent(module *parser.Module, property *qualifiedProperty,
	create bool) (parent *parser.Map, err error) {
	m := &module.Map
	for i, prefix := range property.prefixes() {
		if prop, found := m.GetProperty(prefix); found {
			if mm, ok := prop.Value.(*parser.Map); ok {
				m = mm
			} else {
				// We've found a property in the AST and such property is not of type *parser.Map
				return nil, fmt.Errorf("Expected property %q to be a map, found %s",
					strings.Join(property.prefixes()[:i+1], "."), prop.Value.Type())
			}
		} else if create {
			mm := &parser.Map{}
			m.Properties = append(m.Properties, &parser.Property{Name: prefix, Value: mm})
			m = mm
		} else {
			return nil, nil
		}
	}
	return m, nil
}

type qualifiedProperty struct {
	parts []string
}

func (p *qualifiedProperty) name() string {
	return p.parts[len(p.parts)-1]
}
func (p *qualifiedProperty) prefixes() []string {
	return p.parts[:len(p.parts)-1]
}
func (p *qualifiedProperty) String() string {
	return strings.Join(p.parts, ".")
}

func parseQualifiedProperty(s string) (*qualifiedProperty, error) {
	parts := strings.Split(s, ".")
	if len(parts) == 0 {
		return nil, fmt.Errorf("%q is not a valid property name", s)
	}
	for _, part := range parts {
		if part == "" {
			return nil, fmt.Errorf("%q is not a valid property name", s)
		}
	}
	prop := qualifiedProperty{parts}
	return &prop, nil

}

func parseQualifiedProperties(properties []string) ([]*qualifiedProperty, error) {
	var qualifiedProperties []*qualifiedProperty
	var errs []error
	for _, property := range properties {
		qualifiedProperty, err := parseQualifiedProperty(property)
		if err != nil {
			errs = append(errs, err)
		}
		qualifiedProperties = append(qualifiedProperties, qualifiedProperty)
	}
	if len(errs) > 0 {
		return nil, errors.Join(errs...)
	}
	return qualifiedProperties, nil
}

type Type parser.Type

var (
	List   = Type(parser.ListType)
	String = Type(parser.StringType)
	Bool   = Type(parser.BoolType)
	Int64  = Type(parser.Int64Type)
	Map    = Type(parser.MapType)
)

func (t Type) String() string {
	return parser.Type(t).String()
}