github.com/myhau/pulumi/pkg/v3@v3.70.2-0.20221116134521-f2775972e587/codegen/go/gen_spill.go (about)

     1  package gen
     2  
     3  import (
     4  	"fmt"
     5  
     6  	"github.com/hashicorp/hcl/v2"
     7  	"github.com/pulumi/pulumi/pkg/v3/codegen/hcl2/model"
     8  )
     9  
    10  type spillFunc func(x model.Expression) (string, model.Expression, bool)
    11  
    12  type spillTemp struct {
    13  	Kind     string
    14  	Variable *model.Variable
    15  	Value    model.Expression
    16  }
    17  
    18  type spills struct {
    19  	counts map[string]int
    20  }
    21  
    22  func (s *spills) newTemp(kind string, value model.Expression) *spillTemp {
    23  	i := s.counts[kind]
    24  	s.counts[kind] = i + 1
    25  
    26  	v := &model.Variable{
    27  		Name:         fmt.Sprintf("%s%d", kind, i),
    28  		VariableType: value.Type(),
    29  	}
    30  	return &spillTemp{
    31  		Variable: v,
    32  		Value:    value,
    33  	}
    34  }
    35  
    36  type spiller struct {
    37  	spills *spills
    38  
    39  	temps    []*spillTemp
    40  	spill    spillFunc
    41  	disabled bool
    42  }
    43  
    44  func (s *spiller) preVisit(x model.Expression) (model.Expression, hcl.Diagnostics) {
    45  	_, isfn := x.(*model.AnonymousFunctionExpression)
    46  	if isfn {
    47  		s.disabled = true
    48  	}
    49  	return x, nil
    50  }
    51  
    52  func (s *spiller) postVisit(x model.Expression) (model.Expression, hcl.Diagnostics) {
    53  	_, isfn := x.(*model.AnonymousFunctionExpression)
    54  	if isfn {
    55  		s.disabled = false
    56  	} else if !s.disabled {
    57  		if kind, value, ok := s.spill(x); ok {
    58  			t := s.spills.newTemp(kind, value)
    59  			s.temps = append(s.temps, t)
    60  			return model.VariableReference(t.Variable), nil
    61  		}
    62  	}
    63  	return x, nil
    64  }
    65  
    66  func (g *generator) rewriteSpills(
    67  	x model.Expression, spill spillFunc) (model.Expression, []*spillTemp, hcl.Diagnostics) {
    68  	spiller := &spiller{
    69  		spills: g.spills,
    70  		spill:  spill,
    71  	}
    72  	x, diags := model.VisitExpression(x, spiller.preVisit, spiller.postVisit)
    73  	return x, spiller.temps, diags
    74  }