sanitize.go

Documentation: cuelang.org/go/cue/ast/astutil

     1  // Copyright 2020 CUE Authors
     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  //     http://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 astutil
    16  
    17  import (
    18  	"fmt"
    19  	"math/rand"
    20  	"strings"
    21  
    22  	"cuelang.org/go/cue/ast"
    23  	"cuelang.org/go/cue/errors"
    24  	"cuelang.org/go/cue/token"
    25  )
    26  
    27  // TODO:
    28  // - handle comprehensions
    29  // - change field from foo to "foo" if it isn't referenced, rather than
    30  //   relying on introducing a unique alias.
    31  // - change a predeclared identifier reference to use the __ident form,
    32  //   instead of introducing an alias.
    33  
    34  // Sanitize rewrites File f in place to be well formed after automated
    35  // construction of an AST.
    36  //
    37  // Rewrites:
    38  //   - auto inserts imports associated with Idents
    39  //   - unshadows imports associated with idents
    40  //   - unshadows references for identifiers that were already resolved.
    41  func Sanitize(f *ast.File) error {
    42  	z := &sanitizer{
    43  		file: f,
    44  		rand: rand.New(rand.NewSource(808)),
    45  
    46  		names:      map[string]bool{},
    47  		importMap:  map[string]*ast.ImportSpec{},
    48  		referenced: map[ast.Node]bool{},
    49  		altMap:     map[ast.Node]string{},
    50  	}
    51  
    52  	// Gather all names.
    53  	walk(&scope{
    54  		errFn:   z.errf,
    55  		nameFn:  z.addName,
    56  		identFn: z.markUsed,
    57  	}, f)
    58  	if z.errs != nil {
    59  		return z.errs
    60  	}
    61  
    62  	// Add imports and unshadow.
    63  	s := &scope{
    64  		file:    f,
    65  		errFn:   z.errf,
    66  		identFn: z.handleIdent,
    67  		index:   make(map[string]entry),
    68  	}
    69  	z.fileScope = s
    70  	walk(s, f)
    71  	if z.errs != nil {
    72  		return z.errs
    73  	}
    74  
    75  	z.cleanImports()
    76  
    77  	return z.errs
    78  }
    79  
    80  type sanitizer struct {
    81  	file      *ast.File
    82  	fileScope *scope
    83  
    84  	rand *rand.Rand
    85  
    86  	// names is all used names. Can be used to determine a new unique name.
    87  	names      map[string]bool
    88  	referenced map[ast.Node]bool
    89  
    90  	// altMap defines an alternative name for an existing entry link (a field,
    91  	// alias or let clause). As new names are globally unique, they can be
    92  	// safely reused for any unshadowing.
    93  	altMap    map[ast.Node]string
    94  	importMap map[string]*ast.ImportSpec
    95  
    96  	errs errors.Error
    97  }
    98  
    99  func (z *sanitizer) errf(p token.Pos, msg string, args ...interface{}) {
   100  	z.errs = errors.Append(z.errs, errors.Newf(p, msg, args...))
   101  }
   102  
   103  func (z *sanitizer) addName(name string) {
   104  	z.names[name] = true
   105  }
   106  
   107  func (z *sanitizer) addRename(base string, n ast.Node) (alt string, new bool) {
   108  	if name, ok := z.altMap[n]; ok {
   109  		return name, false
   110  	}
   111  
   112  	name := z.uniqueName(base, false)
   113  	z.altMap[n] = name
   114  	return name, true
   115  }
   116  
   117  func (z *sanitizer) unshadow(parent ast.Node, base string, link ast.Node) string {
   118  	name, ok := z.altMap[link]
   119  	if !ok {
   120  		name = z.uniqueName(base, false)
   121  		z.altMap[link] = name
   122  
   123  		// Insert new let clause at top to refer to a declaration in possible
   124  		// other files.
   125  		let := &ast.LetClause{
   126  			Ident: ast.NewIdent(name),
   127  			Expr:  ast.NewIdent(base),
   128  		}
   129  
   130  		var decls *[]ast.Decl
   131  
   132  		switch x := parent.(type) {
   133  		case *ast.File:
   134  			decls = &x.Decls
   135  		case *ast.StructLit:
   136  			decls = &x.Elts
   137  		default:
   138  			panic(fmt.Sprintf("impossible scope type %T", parent))
   139  		}
   140  
   141  		i := 0
   142  		for ; i < len(*decls); i++ {
   143  			if (*decls)[i] == link {
   144  				break
   145  			}
   146  			if f, ok := (*decls)[i].(*ast.Field); ok && f.Label == link {
   147  				break
   148  			}
   149  		}
   150  
   151  		if i > 0 {
   152  			ast.SetRelPos(let, token.NewSection)
   153  		}
   154  
   155  		a := append((*decls)[:i:i], let)
   156  		*decls = append(a, (*decls)[i:]...)
   157  	}
   158  	return name
   159  }
   160  
   161  func (z *sanitizer) markUsed(s *scope, n *ast.Ident) bool {
   162  	if n.Node != nil {
   163  		return false
   164  	}
   165  	_, _, entry := s.lookup(n.String())
   166  	z.referenced[entry.link] = true
   167  	return true
   168  }
   169  
   170  func (z *sanitizer) cleanImports() {
   171  	z.file.VisitImports(func(d *ast.ImportDecl) {
   172  		k := 0
   173  		for _, s := range d.Specs {
   174  			if _, ok := z.referenced[s]; ok {
   175  				d.Specs[k] = s
   176  				k++
   177  			}
   178  		}
   179  		d.Specs = d.Specs[:k]
   180  	})
   181  }
   182  
   183  func (z *sanitizer) handleIdent(s *scope, n *ast.Ident) bool {
   184  	if n.Node == nil {
   185  		return true
   186  	}
   187  
   188  	_, _, node := s.lookup(n.Name)
   189  	if node.node == nil {
   190  		spec, ok := n.Node.(*ast.ImportSpec)
   191  		if !ok {
   192  			// Clear node. A reference may have been moved to a different
   193  			// file. If not, it should be an error.
   194  			n.Node = nil
   195  			n.Scope = nil
   196  			return false
   197  		}
   198  
   199  		_ = z.addImport(spec)
   200  		info, _ := ParseImportSpec(spec)
   201  		z.fileScope.insert(info.Ident, spec, spec)
   202  		return true
   203  	}
   204  
   205  	if x, ok := n.Node.(*ast.ImportSpec); ok {
   206  		xi, _ := ParseImportSpec(x)
   207  
   208  		if y, ok := node.node.(*ast.ImportSpec); ok {
   209  			yi, _ := ParseImportSpec(y)
   210  			if xi.ID == yi.ID { // name must be identical as a result of lookup.
   211  				z.referenced[y] = true
   212  				n.Node = x
   213  				n.Scope = nil
   214  				return false
   215  			}
   216  		}
   217  
   218  		// Either:
   219  		// - the import is shadowed
   220  		// - an incorrect import is matched
   221  		// In all cases we need to create a new import with a unique name or
   222  		// use a previously created one.
   223  		spec := z.importMap[xi.ID]
   224  		if spec == nil {
   225  			name := z.uniqueName(xi.Ident, false)
   226  			spec = z.addImport(&ast.ImportSpec{
   227  				Name: ast.NewIdent(name),
   228  				Path: x.Path,
   229  			})
   230  			z.importMap[xi.ID] = spec
   231  			z.fileScope.insert(name, spec, spec)
   232  		}
   233  
   234  		info, _ := ParseImportSpec(spec)
   235  		// TODO(apply): replace n itself directly
   236  		n.Name = info.Ident
   237  		n.Node = spec
   238  		n.Scope = nil
   239  		return false
   240  	}
   241  
   242  	if node.node == n.Node {
   243  		return true
   244  	}
   245  
   246  	// n.Node != node and are both not nil and n.Node is not an ImportSpec.
   247  	// This means that either n.Node is illegal or shadowed.
   248  	// Look for the scope in which n.Node is defined and add an alias or let.
   249  
   250  	parent, e, ok := s.resolveScope(n.Name, n.Node)
   251  	if !ok {
   252  		// The node isn't within a legal scope within this file. It may only
   253  		// possibly shadow a value of another file. We add a top-level let
   254  		// clause to refer to this value.
   255  
   256  		// TODO(apply): better would be to have resolve use Apply so that we can replace
   257  		// the entire ast.Ident, rather than modifying it.
   258  		// TODO: resolve to new node or rely on another pass of Resolve?
   259  		n.Name = z.unshadow(z.file, n.Name, n)
   260  		n.Node = nil
   261  		n.Scope = nil
   262  
   263  		return false
   264  	}
   265  
   266  	var name string
   267  	// var isNew bool
   268  	switch x := e.link.(type) {
   269  	case *ast.Field: // referring to regular field.
   270  		name, ok = z.altMap[x]
   271  		if ok {
   272  			break
   273  		}
   274  		// If this field has not alias, introduce one with a unique name.
   275  		// If this has an alias, also introduce a new name. There is a
   276  		// possibility that the alias can be used, but it is easier to just
   277  		// assign a new name, assuming this case is rather rare.
   278  		switch y := x.Label.(type) {
   279  		case *ast.Alias:
   280  			name = z.unshadow(parent, y.Ident.Name, y)
   281  
   282  		case *ast.Ident:
   283  			var isNew bool
   284  			name, isNew = z.addRename(y.Name, x)
   285  			if isNew {
   286  				ident := ast.NewIdent(name)
   287  				// Move formatting and comments from original label to alias
   288  				// identifier.
   289  				CopyMeta(ident, y)
   290  				ast.SetRelPos(y, token.NoRelPos)
   291  				ast.SetComments(y, nil)
   292  				x.Label = &ast.Alias{Ident: ident, Expr: y}
   293  			}
   294  
   295  		default:
   296  			// This is an illegal reference.
   297  			return false
   298  		}
   299  
   300  	case *ast.LetClause:
   301  		name = z.unshadow(parent, x.Ident.Name, x)
   302  
   303  	case *ast.Alias:
   304  		name = z.unshadow(parent, x.Ident.Name, x)
   305  
   306  	default:
   307  		panic(fmt.Sprintf("unexpected link type %T", e.link))
   308  	}
   309  
   310  	// TODO(apply): better would be to have resolve use Apply so that we can replace
   311  	// the entire ast.Ident, rather than modifying it.
   312  	n.Name = name
   313  	n.Node = nil
   314  	n.Scope = nil
   315  
   316  	return true
   317  }
   318  
   319  // uniqueName returns a new name globally unique name of the form
   320  // base_XX ... base_XXXXXXXXXXXXXX or _base or the same pattern with a '_'
   321  // prefix if hidden is true.
   322  //
   323  // It prefers short extensions over large ones, while ensuring the likelihood of
   324  // fast termination is high. There are at least two digits to make it visually
   325  // clearer this concerns a generated number.
   326  func (z *sanitizer) uniqueName(base string, hidden bool) string {
   327  	if hidden && !strings.HasPrefix(base, "_") {
   328  		base = "_" + base
   329  		if !z.names[base] {
   330  			z.names[base] = true
   331  			return base
   332  		}
   333  	}
   334  
   335  	const mask = 0xff_ffff_ffff_ffff // max bits; stay clear of int64 overflow
   336  	const shift = 4                  // rate of growth
   337  	for n := int64(0x10); ; n = int64(mask&((n<<shift)-1)) + 1 {
   338  		num := z.rand.Intn(int(n))
   339  		name := fmt.Sprintf("%s_%01X", base, num)
   340  		if !z.names[name] {
   341  			z.names[name] = true
   342  			return name
   343  		}
   344  	}
   345  }
   346  
   347  func (z *sanitizer) addImport(spec *ast.ImportSpec) *ast.ImportSpec {
   348  	spec = insertImport(&z.file.Decls, spec)
   349  	z.referenced[spec] = true
   350  	return spec
   351  }
   352
View as plain text