template.go

     1  // Copyright 2026 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package specgen
     6  
     7  import (
     8  	"bytes"
     9  	"cmp"
    10  	"fmt"
    11  	"reflect"
    12  	"regexp"
    13  	"simd/archsimd/_gen/specgen/specexpr"
    14  	"slices"
    15  	"strings"
    16  	"text/template"
    17  	"unicode/utf8"
    18  )
    19  
    20  // specTemplate wraps a parsed text/template.Template, or a plain string if no
    21  // template actions were present.
    22  type specTemplate struct {
    23  	raw  string
    24  	tmpl *template.Template
    25  }
    26  
    27  var templateFuncs = template.FuncMap{
    28  	"title": func(v any) string {
    29  		s := fmt.Sprint(v)
    30  		if len(s) == 0 {
    31  			return ""
    32  		}
    33  		return strings.ToTitle(s[:1]) + s[1:]
    34  	},
    35  	"lt": func(a, b any) (bool, error) {
    36  		res, ok := compare(a, b)
    37  		if !ok {
    38  			return false, fmt.Errorf("incomparable types: %T and %T", a, b)
    39  		}
    40  		return res < 0, nil
    41  	},
    42  	"le": func(a, b any) (bool, error) {
    43  		res, ok := compare(a, b)
    44  		if !ok {
    45  			return false, fmt.Errorf("incomparable types: %T and %T", a, b)
    46  		}
    47  		return res <= 0, nil
    48  	},
    49  	"gt": func(a, b any) (bool, error) {
    50  		res, ok := compare(a, b)
    51  		if !ok {
    52  			return false, fmt.Errorf("incomparable types: %T and %T", a, b)
    53  		}
    54  		return res > 0, nil
    55  	},
    56  	"ge": func(a, b any) (bool, error) {
    57  		res, ok := compare(a, b)
    58  		if !ok {
    59  			return false, fmt.Errorf("incomparable types: %T and %T", a, b)
    60  		}
    61  		return res >= 0, nil
    62  	},
    63  	"eq": func(a, b any) bool {
    64  		if res, ok := compare(a, b); ok {
    65  			return res == 0
    66  		}
    67  		return reflect.DeepEqual(a, b)
    68  	},
    69  	"ne": func(a, b any) bool {
    70  		if res, ok := compare(a, b); ok {
    71  			return res != 0
    72  		}
    73  		return !reflect.DeepEqual(a, b)
    74  	},
    75  
    76  	// lanes returns shape.Lanes() if shape is non-scalable; or else "name.Len()"
    77  	"lanes": func(shape, name any) (string, error) {
    78  		s, ok := shape.(specexpr.Vector)
    79  		if !ok {
    80  			return "", fmt.Errorf("want Vector, got %T", shape)
    81  		}
    82  		if !s.Scalable() {
    83  			return s.Elems().String(), nil
    84  		}
    85  		return name.(string) + ".Len()", nil
    86  	},
    87  	"reshapeDiagram": func(x, z any) (string, error) {
    88  		xx, ok := x.(specexpr.Vector)
    89  		if !ok {
    90  			return "", fmt.Errorf("want Vector, got %T", x)
    91  		}
    92  		zz, ok := z.(specexpr.Vector)
    93  		if !ok {
    94  			return "", fmt.Errorf("want Vector, got %T", z)
    95  		}
    96  		if xx.Width != zz.Width {
    97  			return "", fmt.Errorf("Vectors must have equal widths, got %s and %s", xx, zz)
    98  		}
    99  		return reshapeDiagram(xx, zz), nil
   100  	},
   101  }
   102  
   103  func toSpecNum(v any) (specexpr.Num, bool) {
   104  	switch v := v.(type) {
   105  	case specexpr.Num:
   106  		return v, true
   107  	case int:
   108  		return specexpr.Int(v), true
   109  	case int64:
   110  		return specexpr.Int(v), true
   111  	default:
   112  		return nil, false
   113  	}
   114  }
   115  
   116  func compare(a, b any) (int, bool) {
   117  	if na, okA := toSpecNum(a); okA {
   118  		if nb, okB := toSpecNum(b); okB {
   119  			return na.Compare(nb)
   120  		}
   121  	}
   122  	if sa, okA := a.(string); okA {
   123  		if sb, okB := b.(string); okB {
   124  			return cmp.Compare(sa, sb), true
   125  		}
   126  	}
   127  	return 0, false
   128  }
   129  
   130  // reshapeDiagram returns a diagram showing how vector elements are
   131  // reinterpreted by ReshapeToUints. For example, for Uint32x8 -> Uint64x4, it
   132  // returns:
   133  //
   134  //	    x[7]      x[6]    ⋯     x[1]      x[0]
   135  //	| 31 .. 0 | 31 .. 0 | ⋯ | 31 .. 0 | 31 .. 0 |
   136  //	| 63     ....     0 | ⋯ | 63     ....     0 |
   137  //	         z[3]         ⋯          z[0]
   138  func reshapeDiagram(x, z specexpr.Vector) string {
   139  	scalable := x.Scalable() || z.Scalable()
   140  	if !scalable {
   141  		xW, okXW := x.Width.(specexpr.Int)
   142  		zW, okZW := z.Width.(specexpr.Int)
   143  		if !okXW || !okZW || xW != zW {
   144  			return ""
   145  		}
   146  	}
   147  
   148  	// Order x and z by which is smaller and larger
   149  	s, l := x, z
   150  	sName, lName := "x", "z"
   151  	if s.Elem.Bits > l.Elem.Bits {
   152  		s, l = l, s
   153  		sName, lName = lName, sName
   154  	}
   155  
   156  	R := int(l.Elem.Bits / s.Elem.Bits)
   157  
   158  	// Table formatting helpers
   159  	type table []string
   160  	// vstack combines two tables vertically.
   161  	vstack := func(tables ...table) table {
   162  		return slices.Concat(tables...)
   163  	}
   164  	// hstack combines two tables horizontally.
   165  	hstack := func(tables ...table) table {
   166  		out := make(table, len(tables[0]))
   167  		for row := range out {
   168  			for _, table := range tables {
   169  				out[row] += table[row]
   170  			}
   171  		}
   172  		return out
   173  	}
   174  	repeat4 := func(x string) table { return table{x, x, x, x} }
   175  	// format returns left+pad+center+pad+right, such that the length is at least width.
   176  	format := func(left, center, right string, width int) string {
   177  		pad := max(2, width-len(left)-len(center)-len(right)-2)
   178  		// The (pad+1)/2 is because pad == pad/2 + (pad+1)/2 with truncating division.
   179  		return " " + left + strings.Repeat(" ", (pad+1)/2) + center + strings.Repeat(" ", pad/2) + right + " "
   180  	}
   181  
   182  	// formatSmallCells formats the labels and boxes for small vector indexes [from, to].
   183  	sBox := format(fmt.Sprintf("%d", s.Elem.Bits-1), "..", "0", 9)
   184  	formatSmallCells := func(from, to int) table {
   185  		var boxes, labels string
   186  		for sIdx := from; sIdx >= to; sIdx-- {
   187  			labels += format("", fmt.Sprintf("%s[%d]", sName, sIdx), "", len(sBox)) + " "
   188  			boxes += sBox + "|"
   189  		}
   190  		return table{labels, boxes}
   191  	}
   192  
   193  	// formatCol constructs a full table column with small and large cells for a
   194  	// single large vector index.
   195  	formatCol := func(largeIdx int) table {
   196  		high := (largeIdx+1)*R - 1
   197  		low := largeIdx * R
   198  
   199  		// Format the small cells that fan-out to this large cell, This gives us
   200  		// metrics for formatting the large cell.
   201  		var small table
   202  		if R <= 4 {
   203  			// There's enough room to show all of the small elements that form
   204  			// one large element.
   205  			small = formatSmallCells(high, low)
   206  		} else {
   207  			// The fan-out is too high. Show the left and right 2 small cells
   208  			// and elide the middle.
   209  			small = hstack(formatSmallCells(high, high-1), table{" ⋯  ", " ⋯ |"}, formatSmallCells(low+1, low))
   210  		}
   211  
   212  		// Format the large cell.
   213  		width := utf8.RuneCountInString(small[0])
   214  		large := table{
   215  			format(fmt.Sprintf("%d", l.Elem.Bits-1), "....", "0", width-1) + "|",
   216  			format("", fmt.Sprintf("%s[%d]", lName, largeIdx), "", width-1) + " "}
   217  		// Stack small and large cells
   218  		return vstack(small, large)
   219  	}
   220  
   221  	// Format columns for exactly two large vector elements.
   222  	var left, right table
   223  	cellSeps := table{" ", "|", "|", " "}
   224  	prefix := cellSeps
   225  	colSep := table{"", "", "", ""}
   226  	if scalable {
   227  		left, right = formatCol(1), formatCol(0)
   228  		// For scalable vectors, we format large elements 0 and 1 and then show
   229  		// it goes on forever to the left.
   230  		prefix = hstack(repeat4("⋯ "), prefix)
   231  	} else {
   232  		leftIndex := int(l.Elems().(specexpr.Int)) - 1
   233  		left, right = formatCol(leftIndex), formatCol(0)
   234  		if leftIndex > 1 {
   235  			// We're only showing the top and bottom large element, so show that
   236  			// we're eliding the middle elements.
   237  			colSep = hstack(repeat4(" ⋯ "), cellSeps)
   238  		}
   239  	}
   240  
   241  	// Put it all together
   242  	tab := hstack(prefix, left, colSep, right)
   243  
   244  	// We built the smaller elements on top, so flip the table if they should be
   245  	// on bottom.
   246  	if x.Elem.Bits > z.Elem.Bits {
   247  		slices.Reverse(tab)
   248  	}
   249  
   250  	// Finally, emit
   251  	for i := range tab {
   252  		tab[i] = strings.TrimRight(tab[i], " ")
   253  	}
   254  	return "\t  " + strings.Join(tab, "\n\t  ")
   255  }
   256  
   257  // newSpecTemplate parses a spec template string. If tmpl does not contain "{{",
   258  // it is treated as a raw string literal without template overhead.
   259  func newSpecTemplate(tmpl string) (specTemplate, error) {
   260  	if !strings.Contains(tmpl, "{{") {
   261  		return specTemplate{raw: tmpl, tmpl: nil}, nil
   262  	}
   263  
   264  	t, err := template.New("").Option("missingkey=error").Funcs(templateFuncs).Parse(tmpl)
   265  	if err != nil {
   266  		return specTemplate{}, err
   267  	}
   268  	return specTemplate{raw: tmpl, tmpl: t}, nil
   269  }
   270  
   271  func (s *specTemplate) expand(data any) (string, error) {
   272  	if s.tmpl == nil {
   273  		return s.raw, nil
   274  	}
   275  	var buf bytes.Buffer
   276  	if err := s.tmpl.Execute(&buf, data); err != nil {
   277  		return "", err
   278  	}
   279  	return buf.String(), nil
   280  }
   281  
   282  var reMultipleNewlines = regexp.MustCompile(`\n{3,}`)
   283  
   284  func cleanDocNewlines(s string) string {
   285  	s = reMultipleNewlines.ReplaceAllString(s, "\n\n")
   286  	trimmed := strings.TrimRight(s, " \t\n")
   287  	if trimmed == "" {
   288  		return ""
   289  	}
   290  	return trimmed + "\n"
   291  }
   292  
   293  // expandNameAndDoc instantiates the API function name and doc comment for sFn
   294  // using the solved variable bindings in b.
   295  func (sFn *specFunc) expandNameAndDoc(ctx context, b *specexpr.Bindings) (name, doc string) {
   296  	data := make(map[string]any)
   297  	for v, val := range b.All() {
   298  		data[string(v)] = val
   299  		if s, ok := strings.CutPrefix(string(v), "$"); ok {
   300  			data[s] = val
   301  		}
   302  	}
   303  
   304  	var err error
   305  	name, err = sFn.NameTmpl.expand(data)
   306  	if err != nil {
   307  		ctx.errorf("expanding function name template: %s", err)
   308  		name = sFn.Name
   309  	}
   310  
   311  	data["Name"] = name
   312  	doc, err = sFn.Doc.expand(data)
   313  	if err != nil {
   314  		ctx.errorf("expanding doc template: %s", err)
   315  		doc = sFn.Doc.raw
   316  	}
   317  
   318  	doc = cleanDocNewlines(doc)
   319  
   320  	if name != sFn.Name {
   321  		doc = regexp.MustCompile(`\b`+regexp.QuoteMeta(sFn.Name)+`\b`).ReplaceAllLiteralString(doc, name)
   322  	}
   323  
   324  	return name, doc
   325  }
   326
View as plain text