Skip to content

Commit 0ea1051

Browse files
committed
Add count builtin
1 parent 17ab719 commit 0ea1051

File tree

5 files changed

+54
-1
lines changed

5 files changed

+54
-1
lines changed

checker/checker.go

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -492,10 +492,31 @@ func (v *visitor) BuiltinNode(node *ast.BuiltinNode) reflect.Type {
492492
return reflect.ArrayOf(0, closure.Out(0))
493493

494494
}
495-
panic(v.error(node.Arguments[1], "closure should return bool"))
495+
panic(v.error(node.Arguments[1], "closure should has one input and one output param"))
496496
}
497497
panic(v.error(node.Arguments[0], "builtin %v takes only array (got %v)", node.Name, collection))
498498

499+
case "count":
500+
collection := v.visit(node.Arguments[0])
501+
if !isArray(collection) {
502+
panic(v.error(node.Arguments[0], "builtin %v takes only array (got %v)", node.Name, collection))
503+
}
504+
505+
v.collections = append(v.collections, collection)
506+
closure := v.visit(node.Arguments[1])
507+
v.collections = v.collections[:len(v.collections)-1]
508+
509+
if isFunc(closure) &&
510+
closure.NumOut() == 1 &&
511+
closure.NumIn() == 1 && isInterface(closure.In(0)) {
512+
if !isBool(closure.Out(0)) {
513+
panic(v.error(node.Arguments[1], "closure should return boolean (got %v)", closure.Out(0).String()))
514+
}
515+
516+
return integerType
517+
}
518+
panic(v.error(node.Arguments[1], "closure should has one input and one output param"))
519+
499520
default:
500521
panic(v.error(node, "unknown builtin %v", node.Name))
501522
}

checker/checker_test.go

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -157,6 +157,7 @@ func TestCheck(t *testing.T) {
157157
`"a" < "b"`,
158158
"Variadic('', 1, 2) + Variadic('')",
159159
"Foo.Variadic('', 1, 2) + Foo.Variadic('')",
160+
"count(1..30, {# % 3 == 0}) > 0",
160161
}
161162
for _, test := range typeTests {
162163
var err error
@@ -439,6 +440,16 @@ Foo.Variadic('', '')
439440
cannot use string as argument (type int) to call Variadic (1:18)
440441
| Foo.Variadic('', '')
441442
| .................^
443+
444+
count(1, {#})
445+
builtin count takes only array (got int) (1:7)
446+
| count(1, {#})
447+
| ......^
448+
449+
count(ArrayOfInt, {#})
450+
closure should return boolean (got int) (1:19)
451+
| count(ArrayOfInt, {#})
452+
| ..................^
442453
`
443454

444455
func TestCheck_error(t *testing.T) {
@@ -569,6 +580,7 @@ type mockEnv2 struct {
569580
Map map[string]*foo
570581
Any interface{}
571582
ArrayOfAny []interface{}
583+
ArrayOfInt []int
572584
ManOfAny map[string]interface{}
573585
Fn func(bool, int, string, interface{}) string
574586
Bool bool

compiler/compiler.go

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -531,6 +531,21 @@ func (c *compiler) BuiltinNode(node *ast.BuiltinNode) {
531531
c.emit(OpEnd)
532532
c.emit(OpArray)
533533

534+
case "count":
535+
count := c.makeConstant("count")
536+
c.compile(node.Arguments[0])
537+
c.emit(OpBegin)
538+
c.emitPush(0)
539+
c.emit(OpStore, count...)
540+
c.emitLoop(func() {
541+
c.compile(node.Arguments[1])
542+
c.emitCond(func() {
543+
c.emit(OpInc, count...)
544+
})
545+
})
546+
c.emit(OpLoad, count...)
547+
c.emit(OpEnd)
548+
534549
default:
535550
panic(fmt.Sprintf("unknown builtin %v", node.Name))
536551
}

expr_test.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -687,6 +687,10 @@ func TestExpr(t *testing.T) {
687687
`one([1,1,0,1], {# == 0}) and not one([1,0,0,1], {# == 0})`,
688688
true,
689689
},
690+
{
691+
`count(1..30, {# % 3 == 0})`,
692+
10,
693+
},
690694
{
691695
`Now.After(BirthDay)`,
692696
true,

parser/parser.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,7 @@ var builtins = map[string]builtin{
7070
"one": {2},
7171
"filter": {2},
7272
"map": {2},
73+
"count": {2},
7374
}
7475

7576
type parser struct {

0 commit comments

Comments
 (0)