Merge pull request #22 from vhavlena/rule-transformation

Rule transformation

Author: vhavlena

pkg/smt/exprs.go

@@ -6,6 +6,7 @@ import (
 
 	"github.com/open-policy-agent/opa/v1/ast"
 	verr "github.com/vhavlena/verirego/pkg/err"
+	"github.com/vhavlena/verirego/pkg/types"
 )
 
 // ExprTranslator handles the translation of Rego expressions to SMT-LIB format.
@@ -95,6 +96,226 @@ func (et *ExprTranslator) ExprToSmt(expr *ast.Expr) (string, error) {
 	return smtStr, nil
 }
 
+type varDef struct { string; SmtValue }
+
+func (et *ExprTranslator) BodyToSmt(ruleBody *ast.Body) (*SmtProposition,[]varDef,error) {
+	bodySmts := make([]SmtProposition, 0, len(*ruleBody))
+	definedVars := make(map[string]bool, 0)
+	localVarDefs := make([]varDef, 0)
+	for _, expr := range *ruleBody {
+		// single term
+		if term, ok := expr.Terms.(*ast.Term); ok {
+			smtVal,err := et.termToSmtValue(term)
+			if err != nil {
+				return nil, localVarDefs, err
+			}
+			bodySmts = append(bodySmts, *smtVal.Holds())
+			continue
+		}
+
+		// call
+		terms, ok := expr.Terms.([]*ast.Term)
+		if !ok || len(terms) == 0 {
+			return nil, localVarDefs, verr.ErrInvalidEmptyTerm
+		}
+
+		opStr := removeQuotes(terms[0].String())
+		op,err := getOperation(opStr)
+		if err != nil {
+			return nil, localVarDefs, err
+		}
+		
+		arity := op.Decl.Arity()
+		params := len(terms)-1
+		if arity < params {		// the return is a part of the call
+			def, err := et.handleAssigningFunction(opStr, terms)
+			if err != nil {
+				return nil, localVarDefs, err
+			}
+			localVarDefs = append(localVarDefs, *def)
+			definedVars[def.string] = true
+			continue
+		}
+		
+		// we handle ast.Equality separately, because it can be both assignment and comparison, based on the context
+		if op == ast.Equality {
+			if variable,ok := terms[1].Value.(ast.Var); ok {
+				// create variable
+				rhs := terms[2]
+				val,err := et.termToSmtValue(rhs)
+				if err != nil {
+					return nil, localVarDefs, err
+				}
+
+				name := removeQuotes(variable.String())
+				if definedVars[name] != true {
+					localVarDefs = append(localVarDefs, varDef{name, *val})
+					definedVars[name] = true
+				} else {
+					varSmt, err := et.GetVarValue(variable)
+					if err != nil {
+						return nil, nil, err
+					}
+					bodySmts = append(bodySmts, *varSmt.Equals(val))
+				}
+				continue
+			}
+		}
+
+		// Convert all arguments
+		args, err := et.getOperationArgSmts(terms)
+		if err != nil {
+			return nil, localVarDefs, err
+		}
+
+		// Use regoFuncToSmt to get the SMT-LIB string for the operator
+		bodySmt, err := et.getOperationValue(opStr, args, terms)
+		if err != nil {
+			return nil, localVarDefs, err
+		}
+		bodySmts = append(bodySmts, *bodySmt.Holds())
+	}
+
+	bodySmt := And(bodySmts)
+	return bodySmt, localVarDefs, nil
+}
+
+func (et *ExprTranslator) handleAssigningFunction(op string, terms []*ast.Term) (*varDef, error) {
+	if name,ok := terms[len(terms)-1].Value.(ast.Var); ok {	// creating variable
+		// remove assigned variable from call
+		rhs := terms[0:len(terms)-1]
+		args, err := et.getOperationArgSmts(rhs)
+		if err != nil {
+			return nil, err
+		}
+		val, err := et.getOperationValue(op,args,rhs)
+		if err != nil {
+			return nil, err
+		}
+
+		return &varDef{name.String(), *val}, nil
+	}
+	return nil, verr.ErrUnsupportedFunction	// this should be unreachable
+}
+
+func (et *ExprTranslator) getOperationArgSmts(terms []*ast.Term) ([]string, error) {
+	args := make([]string, len(terms)-1)
+	for i := 1; i < len(terms); i++ {
+		s, err := et.termToSmt(terms[i])
+		if err != nil {
+			return nil, err
+		}
+		args[i-1] = s
+	}
+	return args, nil
+}
+
+func (et *ExprTranslator) getOperationValue(op string, args []string, rhs []*ast.Term) (*SmtValue, error) {
+	val,err := et.regoFuncToSmt(op,args,rhs)
+	if err != nil {
+		return nil, err
+	}
+	construct, opType, err := getAtomConstructorForOperation(op)
+	if err != nil {
+		return nil, err
+	}
+	if construct != "" {
+		val = fmt.Sprintf("(%s %s)", construct, val)
+	}
+	return &SmtValue{value: val, depth: 0, atomics: []types.AtomicType{opType,types.AtomicUndef}}, nil	// TODO: user-defined functions
+}
+
+func getOperation(op string) (*ast.Builtin,error) {
+	switch op {
+	case ast.Plus.Name:
+		return ast.Plus,nil
+	case ast.Minus.Name:
+		return ast.Minus,nil
+	case ast.Multiply.Name:
+		return ast.Multiply,nil
+	case ast.Divide.Name:
+		return ast.Divide,nil
+	case ast.Equal.Name:
+		return ast.Equal,nil
+	case ast.Equality.Name:
+		return ast.Equality,nil
+	case ast.Assign.Name:
+		return ast.Assign,nil
+	case ast.GreaterThan.Name:
+		return ast.GreaterThan,nil
+	case ast.GreaterThanEq.Name:
+		return ast.GreaterThanEq,nil
+	case ast.LessThan.Name:
+		return ast.LessThan,nil
+	case ast.LessThanEq.Name:
+		return ast.LessThanEq,nil
+	case ast.Concat.Name:
+		return ast.Concat,nil
+	case ast.Contains.Name:
+		return ast.Contains,nil
+	case ast.StartsWith.Name:
+		return ast.StartsWith,nil
+	case ast.EndsWith.Name:
+		return ast.EndsWith,nil
+	case ast.IndexOf.Name:
+		return ast.IndexOf,nil
+	case ast.Substring.Name:
+		return ast.Substring,nil
+	default:
+		return nil,verr.ErrUnsupportedFunction
+	}
+}
+
+func /*(et *ExprTranslator)*/ getOperationReturnType(opName string) (types.AtomicType,error) {
+	funcMap := map[string]types.AtomicType{
+		ast.Plus.Name:          types.AtomicInt,        // +
+		ast.Minus.Name:         types.AtomicInt,        // -
+		ast.Multiply.Name:      types.AtomicInt,        // *
+		ast.Divide.Name:        types.AtomicInt,        // /
+		ast.Equal.Name:         types.AtomicBoolean,    // ==
+		ast.Equality.Name:      types.AtomicBoolean,    // =
+		ast.Assign.Name:        types.AtomicBoolean,    // :=
+		ast.GreaterThan.Name:   types.AtomicBoolean,    // >
+		ast.GreaterThanEq.Name:	types.AtomicBoolean,    // >=
+		ast.LessThan.Name:      types.AtomicBoolean,    // <
+		ast.LessThanEq.Name:    types.AtomicBoolean,    // <=
+		ast.Concat.Name:        types.AtomicString,     // concat
+		ast.Contains.Name:      types.AtomicBoolean,    // contains
+		ast.StartsWith.Name:    types.AtomicBoolean,    // startswith
+		ast.EndsWith.Name:      types.AtomicBoolean,    // endswith
+		ast.IndexOf.Name:       types.AtomicInt,        // indexof
+		ast.Substring.Name:     types.AtomicString,     // substring
+		// "length" does not exist
+	}
+
+	// TODO: user defined functions
+	if atomicType,found := funcMap[opName]; found {
+		return atomicType,nil
+	}
+	return types.AtomicUndef,verr.ErrUnsupportedFunction
+}
+
+func getAtomConstructorForOperation(op string) (string,types.AtomicType,error) {
+	opType,err := getOperationReturnType(op)
+	if err != nil {
+		return "", "", verr.ErrUnsupportedFunction
+	}
+	return getAtomConstructorFromType(opType),opType,nil
+}
+
+func getAtomConstructorFromType(t types.AtomicType) string {
+	switch t {
+	case types.AtomicString:
+		return "OString"
+	case types.AtomicInt:
+		return "ONumber"
+	case types.AtomicBoolean:
+		return "OBoolean"
+	default:
+		return ""
+	}
+}
+
 // termToSmt converts a Rego AST term to its SMT-LIB string representation.
 //
 // Parameters:
@@ -148,6 +369,112 @@ func (et *ExprTranslator) termToSmt(term *ast.Term) (string, error) {
 	}
 }
 
+func (et *ExprTranslator) termToSmtValue(term *ast.Term) (*SmtValue, error) {
+	switch v := term.Value.(type) {
+	case ast.String:
+		return NewSmtValueFromString(string(v)), nil
+	case ast.Number:
+		if val,ok := v.Int(); ok {
+			return NewSmtValueFromInt(val), nil
+		}
+		return nil,verr.ErrUnsupportedAtomic
+	case ast.Boolean:
+		return NewSmtValueFromBoolean(bool(v)), nil
+	case *ast.Array:
+		return et.arrayToSmt(v)
+	case ast.Object:
+		return et.objectToSmt(v)
+	case ast.Set:
+		// Not directly supported in SMT-LIB, return error
+		return nil, verr.ErrSetConversionNotSupported
+	case ast.Var:
+		return et.GetVarValue(v)
+	case ast.Ref:
+		name := removeQuotes(v[len(v)-1].String())
+		tp, ok := et.TypeTrans.TypeInfo.Types[name]
+		if !ok {
+			return nil, verr.ErrTypeNotFound
+		}
+		return NewSmtValue(name, tp.TypeDepth()), nil
+	case ast.Call:
+		// Handle string functions and other builtins
+		op := removeQuotes(v[0].String())
+		args := make([]string, len(v)-1)
+		for i := 1; i < len(v); i++ {
+			s, err := et.termToSmt(v[i])
+			if err != nil {
+				return nil, err
+			}
+			args[i-1] = s
+		}
+		return et.getOperationValue(op, args, v)
+	default:
+		return nil, fmt.Errorf("%w: %T", verr.ErrUnsupportedTermType, v)
+	}
+}
+
+func (et *ExprTranslator) GetVarValue(v ast.Var) (*SmtValue, error) {
+	return NewSmtValueFromVar(v, et)
+}
+
+func (et *ExprTranslator) arrayToSmt(arr *ast.Array) (*SmtValue, error) {
+	tp, ok := et.TypeTrans.TypeInfo.Types[arr.String()]
+	if !ok {
+		return nil, verr.ErrTypeNotFound
+	}
+
+	depth := tp.TypeDepth()
+	arrSmt := createConstArray("Int", depth)
+
+	for index := range arr.Len() {
+		val := arr.Elem(index)
+		valSmt,err := et.termToSmtValue(val)
+		if err != nil {
+			return nil,err
+		}
+		valSmt = valSmt.WrapToDepth(depth-1)
+		arrSmt = fmt.Sprintf("(store %s %d %s)", arrSmt, index, valSmt.String())
+	}
+
+	return &SmtValue{
+		value: fmt.Sprintf("(OArray%d %s)",depth,arrSmt), 
+		depth: depth,
+	}, nil
+}
+
+func (et *ExprTranslator) objectToSmt(obj ast.Object) (*SmtValue, error) {
+	tp, ok := et.TypeTrans.TypeInfo.Types[obj.String()]
+	if !ok {
+		return nil, verr.ErrTypeNotFound
+	}
+
+	depth := tp.TypeDepth()
+	objSmt := createConstArray("String", depth)
+
+	for _, key := range obj.Keys() {
+		val := obj.Get(key)
+		valSmt,err := et.termToSmtValue(val)
+		if err != nil {
+			return nil,err
+		}
+		valSmt = valSmt.WrapToDepth(depth-1)
+		objSmt = fmt.Sprintf("(store %s %s %s)", objSmt, key.String(), valSmt.String())
+	}
+
+	return &SmtValue{
+		value: fmt.Sprintf("(OObj%d %s)",depth,objSmt),
+		depth: depth,
+	}, nil
+}
+
+func createConstArray(keyType string, depth int) string {
+	undefChild := "OUndef"
+	for d := range depth-1 {
+		undefChild = fmt.Sprintf("(Atom%d %s)",d+1,undefChild)
+	}
+	return fmt.Sprintf("((as const (Array %s OTypeD%d)) %s)",keyType ,depth-1, undefChild)
+}
+
 // regoFuncToSmt converts a Rego function/operator name and its arguments to an SMT-LIB function application string.
 // If the operator is a known built-in, it maps to the corresponding SMT-LIB function. Otherwise, it declares an uninterpreted function
 // with the appropriate parameter and return types (using declareUnintFunc) and returns its application.