func.go

     1  // Copyright 2025 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 asmgen
     6  
     7  import (
     8  	"fmt"
     9  	"slices"
    10  	"strings"
    11  )
    12  
    13  // Note: Exported fields and methods are expected to be used
    14  // by function generators (like the ones in add.go and so on).
    15  // Unexported fields and methods should not be.
    16  
    17  // A Func represents a single assembly function.
    18  type Func struct {
    19  	Name    string
    20  	Asm     *Asm
    21  	inputs  []string       // name of input slices (not beginning with z)
    22  	outputs []string       // names of output slices (beginning with z)
    23  	args    map[string]int // offsets of args, results on stack
    24  }
    25  
    26  // Func starts a new function in the assembly output.
    27  func (a *Asm) Func(decl string) *Func {
    28  	d, ok := strings.CutPrefix(decl, "func ")
    29  	if !ok {
    30  		a.Fatalf("func decl does not begin with 'func '")
    31  	}
    32  	name, d, ok := strings.Cut(d, "(")
    33  	if !ok {
    34  		a.Fatalf("func decl does not have func arg list")
    35  	}
    36  	f := &Func{
    37  		Name: name,
    38  		Asm:  a,
    39  		args: make(map[string]int),
    40  	}
    41  	a.FreeAll()
    42  
    43  	// Parse argument names and types. Quick and dirty.
    44  	// Convert (args) (results) into args, results.
    45  	d = strings.ReplaceAll(d, ") (", ", ")
    46  	d = strings.TrimSuffix(d, ")")
    47  	args := strings.Split(d, ",")
    48  
    49  	// Assign implicit types to all arguments (x, y int -> x int, y int).
    50  	typ := ""
    51  	for i, arg := range slices.Backward(args) {
    52  		arg = strings.TrimSpace(arg)
    53  		if !strings.Contains(arg, " ") {
    54  			if typ == "" {
    55  				a.Fatalf("missing argument type")
    56  			}
    57  			arg += " " + typ
    58  		} else {
    59  			_, typ, _ = strings.Cut(arg, " ")
    60  		}
    61  		args[i] = arg
    62  	}
    63  
    64  	// Record mapping from names to offsets.
    65  	off := 0
    66  	for _, arg := range args {
    67  		name, typ, _ := strings.Cut(arg, " ")
    68  		switch typ {
    69  		default:
    70  			a.Fatalf("unknown type %s", typ)
    71  		case "Word", "uint", "int":
    72  			f.args[name] = off
    73  			off += a.Arch.WordBytes
    74  		case "[]Word":
    75  			if strings.HasPrefix(name, "z") {
    76  				f.outputs = append(f.outputs, name)
    77  			} else {
    78  				f.inputs = append(f.inputs, name)
    79  			}
    80  			f.args[name+"_base"] = off
    81  			f.args[name+"_len"] = off + a.Arch.WordBytes
    82  			f.args[name+"_cap"] = off + 2*a.Arch.WordBytes
    83  			off += 3 * a.Arch.WordBytes
    84  		}
    85  	}
    86  
    87  	a.Printf("\n")
    88  	a.Printf("// %s\n", decl)
    89  	a.Printf("TEXT ·%s(SB), NOSPLIT, $0\n", name)
    90  	if a.Arch.setup != nil {
    91  		a.Arch.setup(f)
    92  	}
    93  	return f
    94  }
    95  
    96  // Arg allocates a new register, copies the named argument (or result) into it,
    97  // and returns that register.
    98  func (f *Func) Arg(name string) Reg {
    99  	return f.ArgHint(name, HintNone)
   100  }
   101  
   102  // ArgHint is like Arg but uses a register allocation hint.
   103  func (f *Func) ArgHint(name string, hint Hint) Reg {
   104  	off, ok := f.args[name]
   105  	if !ok {
   106  		f.Asm.Fatalf("unknown argument %s", name)
   107  	}
   108  	mem := Reg{fmt.Sprintf("%s+%d(FP)", name, off)}
   109  	if hint == HintMemOK && f.Asm.Arch.memOK {
   110  		return mem
   111  	}
   112  	r := f.Asm.RegHint(hint)
   113  	f.Asm.Mov(mem, r)
   114  	return r
   115  }
   116  
   117  // ArgPtr is like Arg but returns a RegPtr.
   118  func (f *Func) ArgPtr(name string) RegPtr {
   119  	return RegPtr(f.Arg(name))
   120  }
   121  
   122  // StoreArg stores src into the named argument (or result).
   123  func (f *Func) StoreArg(src Reg, name string) {
   124  	off, ok := f.args[name]
   125  	if !ok {
   126  		f.Asm.Fatalf("unknown argument %s", name)
   127  	}
   128  	a := f.Asm
   129  	mem := Reg{fmt.Sprintf("%s+%d(FP)", name, off)}
   130  	if src.IsImm() && !a.Arch.memOK {
   131  		r := a.Reg()
   132  		a.Mov(src, r)
   133  		a.Mov(r, mem)
   134  		a.Free(r)
   135  		return
   136  	}
   137  	a.Mov(src, mem)
   138  }
   139
View as plain text