github.com/vescale/zgraph@v0.0.0-20230410094002-959c02d50f95/planner/expr_rewriter.go (about) 1 // Copyright 2022 zGraph Authors. All rights reserved. 2 // 3 // Licensed under the Apache License, Version 2.0 (the "License"); 4 // you may not use this file except in compliance with the License. 5 // You may obtain a copy of the License at 6 // 7 // http://www.apache.org/licenses/LICENSE-2.0 8 // 9 // Unless required by applicable law or agreed to in writing, software 10 // distributed under the License is distributed on an "AS IS" BASIS, 11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 // See the License for the specific language governing permissions and 13 // limitations under the License. 14 15 package planner 16 17 import ( 18 "fmt" 19 20 "github.com/vescale/zgraph/expression" 21 "github.com/vescale/zgraph/parser/ast" 22 ) 23 24 type exprRewriter struct { 25 p LogicalPlan 26 ctxStack []expression.Expression 27 err error 28 } 29 30 func RewriteExpr(expr ast.ExprNode, p LogicalPlan) (expression.Expression, error) { 31 rewriter := &exprRewriter{ 32 p: p, 33 } 34 expr.Accept(rewriter) 35 if rewriter.err != nil { 36 return nil, rewriter.err 37 } 38 return rewriter.ctxStack[0], nil 39 } 40 41 // Enter implements the ast.Visitor interface. 42 func (er *exprRewriter) Enter(n ast.Node) (node ast.Node, skipChildren bool) { 43 switch expr := n.(type) { 44 case *ast.ValueExpr: 45 _ = expr 46 } 47 return n, false 48 } 49 50 // Leave implements the ast.Visitor interface. 51 func (er *exprRewriter) Leave(n ast.Node) (node ast.Node, ok bool) { 52 if er.err != nil { 53 return node, false 54 } 55 56 switch expr := n.(type) { 57 case *ast.ValueExpr: 58 er.ctxStackAppend(&expression.Constant{Value: expr.Datum}) 59 case *ast.BinaryExpr: 60 lExpr := er.ctxStack[er.ctxStackLen()-2] 61 rExpr := er.ctxStack[er.ctxStackLen()-1] 62 er.ctxStackPop(2) 63 binExpr, err := expression.NewBinaryExpr(expr.Op, lExpr, rExpr) 64 if err != nil { 65 er.err = err 66 return n, true 67 } 68 er.ctxStackAppend(binExpr) 69 case *ast.UnaryExpr: 70 input := er.ctxStack[er.ctxStackLen()-1] 71 er.ctxStackPop(1) 72 unaryExpr, err := expression.NewUnaryExpr(expr.Op, input) 73 if err != nil { 74 er.err = err 75 return n, true 76 } 77 er.ctxStackAppend(unaryExpr) 78 case *ast.VariableReference: 79 idx := er.p.Columns().FindColumnIndex(expr.VariableName) 80 if idx == -1 { 81 er.err = fmt.Errorf("unresolved variable %s", expr.VariableName) 82 return n, true 83 } 84 er.ctxStackAppend(&expression.Column{ 85 Index: idx, 86 Name: expr.VariableName, 87 Type: er.p.Columns()[idx].Type, 88 }) 89 case *ast.PropertyAccess: 90 idx := er.p.Columns().FindColumnIndex(expr.VariableName) 91 if idx == -1 { 92 er.err = fmt.Errorf("unresolved variable %s", expr.VariableName) 93 return n, true 94 } 95 col := &expression.Column{ 96 Index: idx, 97 Name: expr.VariableName, 98 Type: er.p.Columns()[idx].Type, 99 } 100 er.ctxStackAppend(&expression.PropertyAccess{ 101 Expr: col, 102 VariableName: expr.VariableName, 103 PropertyName: expr.PropertyName, 104 }) 105 } 106 107 return n, true 108 } 109 110 func (er *exprRewriter) ctxStackLen() int { 111 return len(er.ctxStack) 112 } 113 114 func (er *exprRewriter) ctxStackPop(num int) { 115 l := er.ctxStackLen() 116 er.ctxStack = er.ctxStack[:l-num] 117 } 118 119 func (er *exprRewriter) ctxStackAppend(col expression.Expression) { 120 er.ctxStack = append(er.ctxStack, col) 121 }