hack in segments to the formatter (#269)

* quite ugly but works

* make it work

* fmt

* comment

go.mod

@@ -48,6 +48,7 @@ require (
 	github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
 	github.com/golang/protobuf v1.5.2 // indirect
 	github.com/google/gnostic v0.5.7-v3refs // indirect
+	github.com/google/go-cmp v0.6.0 // indirect
 	github.com/google/go-github/v50 v50.0.0 // indirect
 	github.com/google/go-querystring v1.1.0 // indirect
 	github.com/google/gofuzz v1.1.0 // indirect

go.sum

@@ -115,8 +115,9 @@ github.com/google/go-cmp v0.5.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/
 github.com/google/go-cmp v0.5.1/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/google/go-cmp v0.5.2/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
-github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38=
 github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
+github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
+github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
 github.com/google/go-github/v50 v50.0.0 h1:gdO1AeuSZZK4iYWwVbjni7zg8PIQhp7QfmPunr016Jk=
 github.com/google/go-github/v50 v50.0.0/go.mod h1:Ev4Tre8QoKiolvbpOSG3FIi4Mlon3S2Nt9W5JYqKiwA=
 github.com/google/go-github/v52 v52.0.0 h1:uyGWOY+jMQ8GVGSX8dkSwCzlehU3WfdxQ7GweO/JP7M=

pkg/repo/feature.go

@@ -334,6 +334,33 @@ func (r *repository) Compile(ctx context.Context, req *CompileRequest) ([]*Featu
 	if len(results) < 50 {
 		concurrency = len(results)
 	}
+
+	// This object will hold a map from namespace to named segments
+	// for that ns. For example, if a segment config has two segments: "foo == A, segmentA" and "bar == B, segmentB" in namespace "default"
+	// this map would contain:
+	//
+	// {"default":{"segmentA": "foo == A", "segmentB": "foo == b"}}
+	//
+	nsNameToSegments := make(map[string]map[string]string)
+	for _, fcr := range results {
+		if fcr.FeatureName == "segments" {
+			cf, err := r.CompileFeature(ctx, registry, fcr.NamespaceName, fcr.FeatureName, fcr.NamespaceVersion)
+			if err != nil {
+				fcr.CompilationError = err
+			}
+			rules := make(map[string]string)
+			for _, override := range cf.Feature.Overrides {
+				val, ok := override.Value.(string)
+				if !ok {
+					panic(fmt.Sprintf("should be a string: %+v", override.Value))
+				}
+
+				rules[val] = override.Rule
+			}
+			nsNameToSegments[fcr.NamespaceName] = rules
+		}
+	}
+
 	var wg sync.WaitGroup
 	featureChan := make(chan *FeatureCompilationResult)
 	// start compile workers
@@ -351,7 +378,7 @@ func (r *repository) Compile(ctx context.Context, req *CompileRequest) ([]*Featu
 					}
 					ff := feature.NewFeatureFile(fcr.NamespaceName, fcr.FeatureName)
 					// format starlark file
-					_, fmtDiffExists, err := r.FormatFeature(ctx, &ff, registry, req.DryRun, req.Verbose)
+					_, fmtDiffExists, err := r.FormatFeature(ctx, &ff, registry, req.DryRun, req.Verbose, nsNameToSegments[fcr.NamespaceName])
 					if err != nil {
 						fcr.CompilationError = errors.Wrap(err, "format")
 						return
@@ -365,7 +392,6 @@ func (r *repository) Compile(ctx context.Context, req *CompileRequest) ([]*Featu
 			}
 		}()
 	}
-
 	for _, fcr := range results {
 		featureChan <- fcr
 	}
@@ -546,7 +572,7 @@ func (r *repository) Format(ctx context.Context, verbose bool) error {
 
 		for _, ff := range ffs {
 			ff := ff
-			formatted, _, err := r.FormatFeature(ctx, &ff, nil, false, verbose)
+			formatted, _, err := r.FormatFeature(ctx, &ff, nil, false, verbose, nil)
 			if err != nil {
 				return errors.Wrapf(err, "format config '%s/%s", ff.NamespaceName, ff.Name)
 			}
@@ -558,7 +584,7 @@ func (r *repository) Format(ctx context.Context, verbose bool) error {
 	return nil
 }
 
-func (r *repository) FormatFeature(ctx context.Context, ff *feature.FeatureFile, registry *protoregistry.Types, dryRun, verbose bool) (persisted, diffExists bool, err error) {
+func (r *repository) FormatFeature(ctx context.Context, ff *feature.FeatureFile, registry *protoregistry.Types, dryRun, verbose bool, segments map[string]string) (persisted, diffExists bool, err error) {
 	registry, err = r.registry(ctx, registry)
 	if err != nil {
 		return false, false, err
@@ -574,6 +600,7 @@ func (r *repository) FormatFeature(ctx context.Context, ff *feature.FeatureFile,
 		dryRun,
 		registry,
 		feature.NewNamespaceVersion(nsMD.Version),
+		segments,
 	)
 	// try static formatting
 	persisted, diffExists, err = formatter.StaticFormat(ctx)

pkg/star/formatter.go

@@ -17,6 +17,8 @@ package star
 import (
 	"bytes"
 	"context"
+	"fmt"
+	"strconv"
 
 	"github.com/bazelbuild/buildtools/build"
 	butils "github.com/bazelbuild/buildtools/buildifier/utils"
@@ -39,37 +41,39 @@ type formatter struct {
 	dryRun                bool
 	registry              *protoregistry.Types
 
-	cw fs.ConfigWriter
-	nv feature.NamespaceVersion
+	cw       fs.ConfigWriter
+	nv       feature.NamespaceVersion
+	segments map[string]string
 }
 
-func NewStarFormatter(filePath, featureName string, cw fs.ConfigWriter, dryRun bool, registry *protoregistry.Types, nv feature.NamespaceVersion) Formatter {
+func NewStarFormatter(filePath, featureName string, cw fs.ConfigWriter, dryRun bool, registry *protoregistry.Types, nv feature.NamespaceVersion, segments map[string]string) Formatter {
 	return &formatter{
 		filePath:    filePath,
 		featureName: featureName,
 		cw:          cw,
 		dryRun:      dryRun,
 		registry:    registry,
 		nv:          nv,
+		segments:    segments,
 	}
 }
 
 func (f *formatter) Format(ctx context.Context) (persisted, diffExists bool, err error) {
-	return f.format(ctx, false)
+	return f.format(ctx, false, nil)
 }
 
 func (f *formatter) StaticFormat(ctx context.Context) (persisted, diffExists bool, err error) {
-	return f.format(ctx, true)
+	return f.format(ctx, true, f.segments)
 }
 
-func (f *formatter) format(ctx context.Context, static bool) (persisted, diffExists bool, err error) {
+func (f *formatter) format(ctx context.Context, static bool, segments map[string]string) (persisted, diffExists bool, err error) {
 	data, err := f.cw.GetFileContents(ctx, f.filePath)
 	if err != nil {
 		return false, false, errors.Wrapf(err, "failed to read file %s", f.filePath)
 	}
 	var ndata []byte
 	if static {
-		ndata, err = f.staticFormat(data)
+		ndata, err = f.staticFormat(data, segments)
 	} else {
 		ndata, err = f.bazelFormat(data)
 	}
@@ -97,7 +101,7 @@ func (f *formatter) bazelFormat(data []byte) ([]byte, error) {
 	return build.Format(bfile), nil
 }
 
-func (f *formatter) staticFormat(data []byte) ([]byte, error) {
+func (f *formatter) staticFormat(data []byte, segments map[string]string) ([]byte, error) {
 	// This statically parses the feature
 	// and writes it back to the file without any functional changes.
 	// Doing this ensures that the file is formatted exactly the way that
@@ -106,5 +110,28 @@ func (f *formatter) staticFormat(data []byte) ([]byte, error) {
 	// If static formatting passes, we can skip buildifier formatting.
 	// If static formatting fails, we can show a warning to the user
 	// that the UI will not be able to parse the feature.
-	return static.NewWalker(f.filePath, data, f.registry, f.nv).Format()
+	walker := static.NewWalker(f.filePath, data, f.registry, f.nv)
+	feat, err := walker.Build()
+	if err != nil {
+		return nil, err
+	}
+
+	// This is an intentional hack, that is meant to eventually be replaced
+	// by imports and reusable rules. The actual formatting code too is a mess,
+	// the fact that I have to set RuleAstNew to nil is not great. Due for a cleanup.
+	metadataMap := feat.Feature.Metadata.AsMap()
+	if res, ok := metadataMap["segments"]; ok && res != nil {
+		m, _ := res.(map[string]interface{})
+		for index, segmentName := range m {
+			key, err := strconv.Atoi(index)
+			if err != nil {
+				panic(fmt.Sprintf("invalid key %s", index))
+			}
+			name, _ := segmentName.(string)
+			feat.Feature.Rules.Rules[key].Condition = segments[name]
+			feat.FeatureOld.Tree.Constraints[key].Rule = segments[name]
+			feat.FeatureOld.Tree.Constraints[key].RuleAstNew = nil
+		}
+	}
+	return walker.Mutate(feat)
 }