internal/function: add udf uast_children

Signed-off-by: Manuel Carmona <[email protected]>

Author: mcarmonaa

internal/function/registry.go

@@ -4,11 +4,12 @@ import "gopkg.in/src-d/go-mysql-server.v0/sql"
 
 // Functions for gitbase queries.
 var Functions = sql.Functions{
-	"is_tag":       sql.Function1(NewIsTag),
-	"is_remote":    sql.Function1(NewIsRemote),
-	"language":     sql.FunctionN(NewLanguage),
-	"uast":         sql.FunctionN(NewUAST),
-	"uast_mode":    sql.Function3(NewUASTMode),
-	"uast_xpath":   sql.Function2(NewUASTXPath),
-	"uast_extract": sql.Function2(NewUASTExtract),
+	"is_tag":        sql.Function1(NewIsTag),
+	"is_remote":     sql.Function1(NewIsRemote),
+	"language":      sql.FunctionN(NewLanguage),
+	"uast":          sql.FunctionN(NewUAST),
+	"uast_mode":     sql.Function3(NewUASTMode),
+	"uast_xpath":    sql.Function2(NewUASTXPath),
+	"uast_extract":  sql.Function2(NewUASTExtract),
+	"uast_children": sql.Function1(NewUASTChildren),
 }

internal/function/uast.go

@@ -334,13 +334,12 @@ func (f *UASTXPath) Eval(ctx *sql.Context, row sql.Row) (out interface{}, err er
 			return nil, err
 		}
 
-		for _, n := range ns {
-			data, err := n.Marshal()
-			if err != nil {
-				return nil, err
-			}
-			result = append(result, data)
+		m, err := marshalNodes(ns)
+		if err != nil {
+			return nil, err
 		}
+
+		result = append(result, m...)
 	}
 
 	return result, nil
@@ -359,12 +358,29 @@ func nodesFromBlobArray(data interface{}) ([]*uast.Node, error) {
 		if err := node.Unmarshal(n.([]byte)); err != nil {
 			return nil, err
 		}
+
 		nodes[i] = node
 	}
 
 	return nodes, nil
 }
 
+func marshalNodes(nodes []*uast.Node) ([]interface{}, error) {
+	m := make([]interface{}, 0, len(nodes))
+	for _, n := range nodes {
+		if n != nil {
+			data, err := n.Marshal()
+			if err != nil {
+				return nil, err
+			}
+
+			m = append(m, data)
+		}
+	}
+
+	return m, nil
+}
+
 func (f UASTXPath) String() string {
 	return fmt.Sprintf("uast_xpath(%s, %s)", f.Left, f.Right)
 }
@@ -459,18 +475,7 @@ func getUAST(
 		}
 	}
 
-	var result = make([]interface{}, 0, len(nodes))
-	for _, n := range nodes {
-		if n != nil {
-			node, err := n.Marshal()
-			if err != nil {
-				return nil, err
-			}
-			result = append(result, node)
-		}
-	}
-
-	return result, nil
+	return marshalNodes(nodes)
 }
 
 // UASTExtract extracts keys from an UAST.
@@ -586,3 +591,73 @@ func (u *UASTExtract) TransformUp(f sql.TransformExprFunc) (sql.Expression, erro
 
 	return f(NewUASTExtract(left, rigth))
 }
+
+// UASTChildren returns children from UAST nodes.
+type UASTChildren struct {
+	expression.UnaryExpression
+}
+
+// NewUASTChildren creates a new UASTExtract UDF.
+func NewUASTChildren(uast sql.Expression) sql.Expression {
+	return &UASTChildren{expression.UnaryExpression{Child: uast}}
+}
+
+// String implements the fmt.Stringer interface.
+func (u *UASTChildren) String() string {
+	return fmt.Sprintf("uast_children(%s)", u.Child)
+}
+
+// Type implements the sql.Expression interface.
+func (u *UASTChildren) Type() sql.Type {
+	return sql.Array(sql.Blob)
+}
+
+// TransformUp implements the sql.Expression interface.
+func (u *UASTChildren) TransformUp(f sql.TransformExprFunc) (sql.Expression, error) {
+	child, err := u.Child.TransformUp(f)
+	if err != nil {
+		return nil, err
+	}
+
+	return f(NewUASTChildren(child))
+}
+
+// Eval implements the sql.Expression interface.
+func (u *UASTChildren) Eval(ctx *sql.Context, row sql.Row) (out interface{}, err error) {
+	defer func() {
+		if r := recover(); r != nil {
+			err = fmt.Errorf("uast: unknown error: %s", r)
+		}
+	}()
+
+	span, ctx := ctx.Span("gitbase.UASTChildren")
+	defer span.Finish()
+
+	child, err := u.Child.Eval(ctx, row)
+	if err != nil {
+		return nil, err
+	}
+
+	if child == nil {
+		return nil, nil
+	}
+
+	nodes, err := nodesFromBlobArray(child)
+	if err != nil {
+		return nil, err
+	}
+
+	children := flattenChildren(nodes)
+	return marshalNodes(children)
+}
+
+func flattenChildren(nodes []*uast.Node) []*uast.Node {
+	children := []*uast.Node{}
+	for _, n := range nodes {
+		if len(n.Children) > 0 {
+			children = append(children, n.Children...)
+		}
+	}
+
+	return children
+}

internal/function/uast_test.go

@@ -281,6 +281,45 @@ func TestUASTExtract(t *testing.T) {
 	}
 }
 
+func TestUASTChildren(t *testing.T) {
+	var require = require.New(t)
+
+	ctx, cleanup := setup(t)
+	defer cleanup()
+
+	uasts, _ := bblfshFixtures(t, ctx)
+	modes := []string{"semantic", "annotated", "native"}
+
+	for _, mode := range modes {
+		root, ok := uasts[mode]
+		require.True(ok)
+
+		nodes, err := nodesFromBlobArray(root)
+		require.NoError(err)
+		require.Len(nodes, 1)
+		expected := nodes[0].Children
+
+		row := sql.NewRow(root)
+
+		fn := NewUASTChildren(
+			expression.NewGetField(0, sql.Array(sql.Blob), "", false),
+		)
+
+		children, err := fn.Eval(ctx, row)
+		require.NoError(err)
+
+		nodes, err = nodesFromBlobArray(children)
+		require.NoError(err)
+		require.Len(nodes, len(expected))
+		for i, n := range nodes {
+			require.Equal(
+				n.InternalType,
+				expected[i].InternalType,
+			)
+		}
+	}
+}
+
 func assertUASTBlobs(t *testing.T, a, b interface{}) {
 	t.Helper()
 	require := require.New(t)