github.com/wzzhu/tensor@v0.9.24/genlib2/api_unary.go (about)

     1  package main
     2  
     3  import (
     4  	"io"
     5  	"text/template"
     6  )
     7  
     8  type APIUnary struct {
     9  	UnaryOp
    10  }
    11  
    12  func (fn *APIUnary) Signature() *Signature {
    13  	var paramNames []string
    14  	var paramTemplates []*template.Template
    15  	switch {
    16  	case fn.UnaryOp.Name() == "Clamp":
    17  		paramNames = []string{"a", "min", "max", "opts"}
    18  		paramTemplates = []*template.Template{tensorType, interfaceType, interfaceType, splatFuncOptType}
    19  	default:
    20  		paramNames = []string{"a", "opts"}
    21  		paramTemplates = []*template.Template{tensorType, splatFuncOptType}
    22  	}
    23  	return &Signature{
    24  		Name:            fn.Name(),
    25  		NameTemplate:    plainName,
    26  		ParamNames:      paramNames,
    27  		ParamTemplates:  paramTemplates,
    28  		RetVals:         []string{"retVal"},
    29  		RetValTemplates: []*template.Template{tensorType},
    30  		Err:             true,
    31  	}
    32  }
    33  
    34  func (fn *APIUnary) WriteBody(w io.Writer) {
    35  	body := `e := a.Engine()
    36  	if {{interfaceName .Name | lower}}, ok := e.({{interfaceName .Name}}); ok {
    37  		{{if eq .Name "Clamp" -}}
    38  		return clamper.Clamp(a, min, max, opts...)
    39  		{{else -}}
    40  		return {{interfaceName .Name|lower}}.{{.Name}}(a, opts...)
    41  		{{end -}}
    42  	}
    43  	err = errors.Errorf("Engine does not perform {{.Name}}")
    44  	return
    45  	`
    46  
    47  	T := template.Must(template.New("body").Funcs(funcs).Parse(body))
    48  	T.Execute(w, fn)
    49  }
    50  
    51  func (fn *APIUnary) Write(w io.Writer) {
    52  	w.Write([]byte("func "))
    53  	sig := fn.Signature()
    54  	sig.Write(w)
    55  	w.Write([]byte("{ \n"))
    56  	fn.WriteBody(w)
    57  	w.Write([]byte("}\n\n"))
    58  }
    59  
    60  func generateUncondUnaryAPI(f io.Writer, kinds Kinds) {
    61  	var unaries []*APIUnary
    62  	for _, u := range unconditionalUnaries {
    63  		fn := &APIUnary{
    64  			UnaryOp: u,
    65  		}
    66  		unaries = append(unaries, fn)
    67  	}
    68  	for _, u := range unaries {
    69  		u.Write(f)
    70  	}
    71  }
    72  
    73  func generateCondUnaryAPI(f io.Writer, kinds Kinds) {
    74  	var unaries []*APIUnary
    75  	for _, u := range conditionalUnaries {
    76  		fn := &APIUnary{
    77  			UnaryOp: u,
    78  		}
    79  		unaries = append(unaries, fn)
    80  	}
    81  	for _, u := range unaries {
    82  		u.Write(f)
    83  	}
    84  }
    85  
    86  func generateSpecialUnaryAPI(f io.Writer, kinds Kinds) {
    87  	var unaries []*APIUnary
    88  
    89  	for _, u := range specialUnaries {
    90  		fn := &APIUnary{
    91  			UnaryOp: u,
    92  		}
    93  		unaries = append(unaries, fn)
    94  	}
    95  	for _, u := range unaries {
    96  		u.Write(f)
    97  	}
    98  }