vitess.io/vitess@v0.16.2/go/test/dbg/dbg.go (about)

     1  /*
     2  Copyright 2023 The Vitess Authors
     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  	http://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 dbg
    18  
    19  import (
    20  	"bytes"
    21  	"fmt"
    22  	"go/ast"
    23  	"go/format"
    24  	"go/parser"
    25  	"go/token"
    26  	"os"
    27  	"path"
    28  	"runtime"
    29  	"strings"
    30  	"sync"
    31  
    32  	"github.com/kr/pretty"
    33  	"github.com/kr/text"
    34  )
    35  
    36  type params struct {
    37  	pos    token.Position
    38  	fn     string
    39  	params []string
    40  }
    41  
    42  func (p *params) Position() string {
    43  	if p == nil {
    44  		return "<unknown>"
    45  	}
    46  	return p.pos.String()
    47  }
    48  
    49  func (p *params) ShortPosition() string {
    50  	if p == nil {
    51  		return "<unknown>"
    52  	}
    53  	return fmt.Sprintf("%s:%d", path.Base(p.pos.Filename), p.pos.Line)
    54  }
    55  
    56  func (p *params) Arg(n int) string {
    57  	if p == nil || n >= len(p.params) {
    58  		return "arg"
    59  	}
    60  	return p.params[n]
    61  }
    62  
    63  func (p *params) Fn() string {
    64  	if p == nil || p.fn == "" {
    65  		return "?"
    66  	}
    67  	return p.fn
    68  }
    69  
    70  type file struct {
    71  	once  sync.Once
    72  	fset  *token.FileSet
    73  	path  string
    74  	calls map[int]*params
    75  }
    76  
    77  type cache struct {
    78  	mu     sync.Mutex
    79  	parsed map[string]*file
    80  	fset   *token.FileSet
    81  }
    82  
    83  func (f *file) parse() {
    84  	a, err := parser.ParseFile(f.fset, f.path, nil, 0)
    85  	if err != nil {
    86  		_, _ = fmt.Fprintf(os.Stderr, "[dbg] failed to parse %q: %v\n", f.path, err)
    87  		return
    88  	}
    89  
    90  	f.calls = map[int]*params{}
    91  
    92  	var curfn string
    93  	ast.Inspect(a, func(node ast.Node) bool {
    94  		switch n := node.(type) {
    95  		case *ast.FuncDecl:
    96  			var buf strings.Builder
    97  			if n.Recv != nil && len(n.Recv.List) == 1 {
    98  				buf.WriteByte('(')
    99  				_ = format.Node(&buf, f.fset, n.Recv.List[0].Type)
   100  				buf.WriteString(").")
   101  			}
   102  			buf.WriteString(n.Name.String())
   103  			curfn = buf.String()
   104  
   105  		case *ast.CallExpr:
   106  			if sel, ok := n.Fun.(*ast.SelectorExpr); ok {
   107  				if pkg, ok := sel.X.(*ast.Ident); ok {
   108  					if pkg.Name == "dbg" && (sel.Sel.Name == "P" || sel.Sel.Name == "V") {
   109  						var p = params{
   110  							pos: f.fset.Position(n.Pos()),
   111  							fn:  curfn,
   112  						}
   113  
   114  						for _, arg := range n.Args {
   115  							var buf strings.Builder
   116  							_ = format.Node(&buf, f.fset, arg)
   117  							p.params = append(p.params, buf.String())
   118  						}
   119  
   120  						f.calls[p.pos.Line] = &p
   121  						return false
   122  					}
   123  				}
   124  			}
   125  		}
   126  		return true
   127  	})
   128  }
   129  
   130  func (f *file) resolve(lineno int) *params {
   131  	f.once.Do(f.parse)
   132  	return f.calls[lineno]
   133  }
   134  
   135  func (c *cache) resolve(filename string, lineno int) *params {
   136  	var f *file
   137  
   138  	c.mu.Lock()
   139  	f = c.parsed[filename]
   140  	if f == nil {
   141  		f = &file{fset: c.fset, path: filename}
   142  		c.parsed[filename] = f
   143  	}
   144  	c.mu.Unlock()
   145  
   146  	return f.resolve(lineno)
   147  }
   148  
   149  var defaultCache = cache{
   150  	fset:   token.NewFileSet(),
   151  	parsed: map[string]*file{},
   152  }
   153  
   154  // V prints the given argument in compact debug form and returns it unchanged
   155  func V[Val any](v Val) Val {
   156  	var p *params
   157  	if _, f, lineno, ok := runtime.Caller(1); ok {
   158  		p = defaultCache.resolve(f, lineno)
   159  	}
   160  	_, _ = fmt.Fprintf(os.Stdout, "[%s]: %s = %# v\n", p.ShortPosition(), p.Arg(0), pretty.Formatter(v))
   161  	return v
   162  }
   163  
   164  // P prints all the arguments passed to the function in verbose debug form
   165  func P(vals ...any) {
   166  	var p *params
   167  	if _, f, lineno, ok := runtime.Caller(1); ok {
   168  		p = defaultCache.resolve(f, lineno)
   169  	}
   170  
   171  	var buf bytes.Buffer
   172  	_, _ = fmt.Fprintf(&buf, "%s @ %s\n", p.Position(), p.Fn())
   173  	for i, v := range vals {
   174  		indent, _ := fmt.Fprintf(&buf, "    [%d] %s = ", i, p.Arg(i))
   175  
   176  		w := text.NewIndentWriter(&buf, nil, bytes.Repeat([]byte{' '}, indent))
   177  		fmt.Fprintf(w, "%# v\n", pretty.Formatter(v))
   178  	}
   179  	_, _ = buf.WriteTo(os.Stdout)
   180  }