Replace iota variable in constant

tygo/iota.go

@@ -1,37 +1,17 @@
 package tygo
 
 import (
-	"log"
+	"regexp"
 	"strconv"
 	"strings"
 )
 
-func isProbablyIotaType(groupType string) bool {
-	groupType = strings.Trim(groupType, "()")
-	return groupType == "iota" || strings.HasPrefix(groupType, "iota +") || strings.HasSuffix(groupType, "+ iota")
-}
-
-// Note: this could be done so much more elegantly, but this probably covers 99.9% of iota usecases
-func basicIotaOffsetValueParse(groupType string) int {
-	if !isProbablyIotaType(groupType) {
-		panic("can't parse non-iota type")
-	}
-	groupType = strings.Trim(groupType, "()")
-	if groupType == "iota" {
-		return 0
-	}
-	parts := strings.Split(groupType, " + ")
+var iotaRegexp = regexp.MustCompile(`\biota\b`)
 
-	var numPart string
-	if parts[0] == "iota" {
-		numPart = parts[1]
-	} else {
-		numPart = parts[0]
-	}
+func isProbablyIotaType(valueString string) bool {
+	return !strings.ContainsAny(valueString, "\"'`") && iotaRegexp.MatchString(valueString)
+}
 
-	addValue, err := strconv.ParseInt(numPart, 10, 64)
-	if err != nil {
-		log.Panicf("Failed to guesstimate initial iota value for \"%s\": %v", groupType, err)
-	}
-	return int(addValue)
+func replaceIotaValue(valueString string, iotaValue int) string {
+	return iotaRegexp.ReplaceAllLiteralString(valueString, strconv.Itoa(iotaValue))
 }

tygo/write_toplevel.go

@@ -15,7 +15,6 @@ type groupContext struct {
 	groupValue           string
 	groupType            string
 	iotaValue            int
-	iotaOffset           int
 }
 
 func (g *PackageGenerator) writeGroupDecl(s *strings.Builder, decl *ast.GenDecl) {
@@ -202,24 +201,14 @@ func (g *PackageGenerator) writeValueSpec(
 			val := vs.Values[i]
 			tempSB := &strings.Builder{}
 			g.writeType(tempSB, val, 0, true)
-			valueString := tempSB.String()
-
-			if isProbablyIotaType(valueString) {
-				group.iotaOffset = basicIotaOffsetValueParse(valueString)
-				group.groupValue = "iota"
-				valueString = fmt.Sprint(group.iotaValue + group.iotaOffset)
-			} else {
-				group.groupValue = valueString
-			}
-			s.WriteString(valueString)
+			group.groupValue = tempSB.String()
+		}
 
-		} else { // We must use the previous value or +1 in case of iota
-			valueString := group.groupValue
-			if group.groupValue == "iota" {
-				valueString = fmt.Sprint(group.iotaValue + group.iotaOffset)
-			}
-			s.WriteString(valueString)
+		valueString := group.groupValue
+		if isProbablyIotaType(valueString) {
+			valueString = replaceIotaValue(valueString, group.iotaValue)
 		}
+		s.WriteString(valueString)
 
 		s.WriteByte(';')