Implement cap() builtin for type params.

nevkontakte · nevkontakte · commit bfa92fe08542 · 2023-11-12T17:46:23.000Z
Similar to len(), to maintain performance when used in loops I've kept
compile-time specializations for when the exact type is known.
Otherwise, a type-specific implementation is called at runtime.
diff --git a/compiler/expressions.go b/compiler/expressions.go
@@ -1084,9 +1084,13 @@ func (fc *funcContext) translateBuiltin(name string, sig *types.Signature, args
 			return fc.formatExpr("%e.$capacity", args[0])
 		case *types.Pointer:
 			return fc.formatExpr("(%e, %d)", args[0], argType.Elem().(*types.Array).Len())
-		// capacity of array is constant
+		case *types.Array:
+			// This should never happen™
+			panic(fmt.Errorf("array capacity should have been inlined as constant"))
+		case *types.Interface: // *types.TypeParam has interface as underlying type.
+			return fc.formatExpr("%s.$cap(%e)", fc.typeName(fc.pkgCtx.TypeOf(args[0])), args[0])
 		default:
-			panic(fmt.Sprintf("Unhandled cap type: %T\n", argType))
+			panic(fmt.Errorf("unhandled cap type: %T", argType))
 		}
 	case "panic":
 		return fc.formatExpr("$panic(%s)", fc.translateImplicitConversion(args[0], types.NewInterface(nil, nil)))
diff --git a/compiler/prelude/types.js b/compiler/prelude/types.js
@@ -165,6 +165,7 @@ var $newType = (size, kind, string, named, pkg, exported, constructor) => {
                 Object.defineProperty(typ.ptr.nil, "nilCheck", { get: $throwNilPointerError });
             };
             typ.$len = (v) => typ.len;
+            typ.$cap = (v) => typ.len;
             break;
 
         case $kindChan:
@@ -179,6 +180,7 @@ var $newType = (size, kind, string, named, pkg, exported, constructor) => {
             };
             typ.$make = (bufsize) => new $Chan(typ.elem, bufsize ? bufsize : 0);
             typ.$len = (v) => v.$buffer.length;
+            typ.$cap = (v) => v.$capacity;
             break;
 
         case $kindFunc:
@@ -241,6 +243,7 @@ var $newType = (size, kind, string, named, pkg, exported, constructor) => {
                 typ.nil = new typ($throwNilPointerError, $throwNilPointerError);
             };
             typ.$len = (v) => typ.elem.$len(typ.wrapped ? v : v.$get());
+            typ.$cap = (v) => typ.elem.$cap(typ.wrapped ? v : v.$get());
             break;
 
         case $kindSlice:
@@ -263,6 +266,7 @@ var $newType = (size, kind, string, named, pkg, exported, constructor) => {
             };
             typ.$make = (size, capacity) => $makeSlice(typ, size, capacity);
             typ.$len = (v) => v.$length;
+            typ.$cap = (v) => v.$capacity;
             break;
 
         case $kindStruct:
diff --git a/tests/typeparams/builtins_test.go b/tests/typeparams/builtins_test.go
@@ -124,3 +124,42 @@ func TestLen(t *testing.T) {
 		})
 	}
 }
+
+func _cap[T []int | *[3]int | [3]int | chan int](x T) int {
+	return cap(x)
+}
+
+func TestCap(t *testing.T) {
+	ch := make(chan int, 2)
+	ch <- 1
+
+	tests := []struct {
+		desc string
+		got  int
+		want int
+	}{{
+		desc: "[]int",
+		got:  _cap([]int{1, 2, 3}),
+		want: 3,
+	}, {
+		desc: "*[3]int",
+		got:  _cap(&[3]int{1}),
+		want: 3,
+	}, {
+		desc: "[3]int",
+		got:  _cap([3]int{1}),
+		want: 3,
+	}, {
+		desc: "chan int",
+		got:  _cap(ch),
+		want: 2,
+	}}
+
+	for _, test := range tests {
+		t.Run(test.desc, func(t *testing.T) {
+			if test.got != test.want {
+				t.Errorf("Got: len() = %d. Want: %d.", test.got, test.want)
+			}
+		})
+	}
+}
