github.com/rajeev159/opa@v0.45.0/topdown/numbers.go (about) 1 // Copyright 2020 The OPA Authors. All rights reserved. 2 // Use of this source code is governed by an Apache2 3 // license that can be found in the LICENSE file. 4 5 package topdown 6 7 import ( 8 "fmt" 9 "math/big" 10 11 "github.com/open-policy-agent/opa/ast" 12 "github.com/open-policy-agent/opa/topdown/builtins" 13 ) 14 15 type randIntCachingKey string 16 17 var one = big.NewInt(1) 18 19 func builtinNumbersRange(bctx BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error { 20 21 x, err := builtins.BigIntOperand(operands[0].Value, 1) 22 if err != nil { 23 return err 24 } 25 26 y, err := builtins.BigIntOperand(operands[1].Value, 2) 27 if err != nil { 28 return err 29 } 30 31 result := ast.NewArray() 32 cmp := x.Cmp(y) 33 haltErr := Halt{ 34 Err: &Error{ 35 Code: CancelErr, 36 Message: "numbers.range: timed out before generating all numbers in range", 37 }, 38 } 39 40 if cmp <= 0 { 41 for i := new(big.Int).Set(x); i.Cmp(y) <= 0; i = i.Add(i, one) { 42 if bctx.Cancel != nil && bctx.Cancel.Cancelled() { 43 return haltErr 44 } 45 result = result.Append(ast.NewTerm(builtins.IntToNumber(i))) 46 } 47 } else { 48 for i := new(big.Int).Set(x); i.Cmp(y) >= 0; i = i.Sub(i, one) { 49 if bctx.Cancel != nil && bctx.Cancel.Cancelled() { 50 return haltErr 51 } 52 result = result.Append(ast.NewTerm(builtins.IntToNumber(i))) 53 } 54 } 55 56 return iter(ast.NewTerm(result)) 57 } 58 59 func builtinRandIntn(bctx BuiltinContext, args []*ast.Term, iter func(*ast.Term) error) error { 60 61 strOp, err := builtins.StringOperand(args[0].Value, 1) 62 if err != nil { 63 return err 64 65 } 66 67 n, err := builtins.IntOperand(args[1].Value, 2) 68 if err != nil { 69 return err 70 } 71 72 if n == 0 { 73 return iter(ast.IntNumberTerm(0)) 74 } 75 76 if n < 0 { 77 n = -n 78 } 79 80 var key = randIntCachingKey(fmt.Sprintf("%s-%d", strOp, n)) 81 82 if val, ok := bctx.Cache.Get(key); ok { 83 return iter(val.(*ast.Term)) 84 } 85 86 r, err := bctx.Rand() 87 if err != nil { 88 return err 89 } 90 result := ast.IntNumberTerm(r.Intn(n)) 91 bctx.Cache.Put(key, result) 92 93 return iter(result) 94 } 95 96 func init() { 97 RegisterBuiltinFunc(ast.NumbersRange.Name, builtinNumbersRange) 98 RegisterBuiltinFunc(ast.RandIntn.Name, builtinRandIntn) 99 }