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  }