1// Copyright 2021 The Prometheus Authors
2// Licensed under the Apache License, Version 2.0 (the "License");
3// you may not use this file except in compliance with the License.
4// You may obtain a copy of the License at
5//
6// http://www.apache.org/licenses/LICENSE-2.0
7//
8// Unless required by applicable law or agreed to in writing, software
9// distributed under the License is distributed on an "AS IS" BASIS,
10// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11// See the License for the specific language governing permissions and
12// limitations under the License.
13
14import { SyntaxNode } from '@lezer/common';
15import {
16  AggregateExpr,
17  BinaryExpr,
18  Expr,
19  FunctionCall,
20  MatrixSelector,
21  NumberLiteral,
22  OffsetExpr,
23  ParenExpr,
24  StepInvariantExpr,
25  StringLiteral,
26  SubqueryExpr,
27  UnaryExpr,
28  VectorSelector,
29} from '../grammar/parser.terms';
30import { walkThrough } from './path-finder';
31import { getFunction, ValueType } from '../types';
32
33// Based on https://github.com/prometheus/prometheus/blob/d668a7efe3107dbdcc67bf4e9f12430ed8e2b396/promql/parser/ast.go#L191
34export function getType(node: SyntaxNode | null): ValueType {
35  if (!node) {
36    return ValueType.none;
37  }
38  switch (node.type.id) {
39    case Expr:
40      return getType(node.firstChild);
41    case AggregateExpr:
42      return ValueType.vector;
43    case VectorSelector:
44      return ValueType.vector;
45    case OffsetExpr:
46      return getType(node.firstChild);
47    case StringLiteral:
48      return ValueType.string;
49    case NumberLiteral:
50      return ValueType.scalar;
51    case MatrixSelector:
52      return ValueType.matrix;
53    case SubqueryExpr:
54      return ValueType.matrix;
55    case ParenExpr:
56      return getType(walkThrough(node, Expr));
57    case UnaryExpr:
58      return getType(walkThrough(node, Expr));
59    case BinaryExpr:
60      const lt = getType(node.firstChild);
61      const rt = getType(node.lastChild);
62      if (lt === ValueType.scalar && rt === ValueType.scalar) {
63        return ValueType.scalar;
64      }
65      return ValueType.vector;
66    case FunctionCall:
67      const funcNode = node.firstChild?.firstChild;
68      if (!funcNode) {
69        return ValueType.none;
70      }
71      return getFunction(funcNode.type.id).returnType;
72    case StepInvariantExpr:
73      return getType(walkThrough(node, Expr));
74    default:
75      return ValueType.none;
76  }
77}
78