hrw4u: add an AST round-trip test over the whole corpus

tools/hrw4u/src/ast_nodes.py

@@ -25,16 +25,22 @@
     "IdentValue",
     "IPValue",
     "ParamRef",
+    "BoolValue",
+    "NumberValue",
     "RegexValue",
+    "SetValue",
+    "IpRangeValue",
     "ValueExpr",
     "Node",
     "Target",
     "Assignment",
     "FunctionCall",
     "Break",
+    "Comment",
     "Comparison",
     "LogicalOp",
     "NotOp",
+    "Group",
     "BoolLiteral",
     "IdentCondition",
     "ElifBranch",
@@ -72,12 +78,32 @@ class ParamRef:
     raw: str
 
 
+@dataclass(frozen=True, kw_only=True)
+class BoolValue:
+    raw: str  # source spelling, e.g. "TRUE", "true", "TRue"
+
+
+@dataclass(frozen=True, kw_only=True)
+class NumberValue:
+    raw: str  # source spelling: header_rewrite echoes it, so 007 is three bytes and not 7
+
+
 @dataclass(frozen=True, kw_only=True)
 class RegexValue:
     raw: str
 
 
-ValueExpr = Union[LiteralStringValue, IdentValue, IPValue, ParamRef, int, bool, tuple[IPValue, ...]]
+@dataclass(frozen=True, kw_only=True)
+class SetValue:
+    raw: str  # bracket-stripped source text, quoting preserved
+
+
+@dataclass(frozen=True, kw_only=True)
+class IpRangeValue:
+    raw: str  # verbatim source text
+
+
+ValueExpr = Union[LiteralStringValue, IdentValue, IPValue, ParamRef, BoolValue, NumberValue, IpRangeValue]
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -119,11 +145,16 @@ class Break(Node):
     pass
 
 
+@dataclass(frozen=True, kw_only=True)
+class Comment(Node):
+    text: str
+
+
 @dataclass(frozen=True, kw_only=True)
 class Comparison(Node):
     left: IdentValue | FunctionCall
     operator: str  # "==", "!=", ">", "<", "~", "!~", "in", "!in"
-    right: ValueExpr | RegexValue | tuple[ValueExpr, ...]
+    right: ValueExpr | RegexValue | SetValue | IpRangeValue
     modifiers: tuple[str, ...]
 
 
@@ -139,6 +170,11 @@ class NotOp(Node):
     operand: ConditionExpr
 
 
+@dataclass(frozen=True, kw_only=True)
+class Group(Node):
+    inner: ConditionExpr
+
+
 @dataclass(frozen=True, kw_only=True)
 class BoolLiteral(Node):
     value: bool
@@ -161,6 +197,7 @@ class IfBlock(Node):
     body: tuple[BodyNode, ...]
     elif_branches: tuple[ElifBranch, ...]
     else_body: tuple[BodyNode, ...]
+    has_else: bool  # an empty else body is not the same as no else clause
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -185,7 +222,7 @@ class VarDecl(Node):
 @dataclass(frozen=True, kw_only=True)
 class VarSection(Node):
     scope: str
-    declarations: tuple[VarDecl, ...]
+    declarations: tuple[VarDecl | Comment, ...]
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -206,6 +243,6 @@ class HRW4UAST:
 
 
 # Type aliases: must follow all class definitions (evaluated at runtime).
-ConditionExpr = Union[Comparison, LogicalOp, NotOp, BoolLiteral, IdentCondition, FunctionCall]
-BodyNode = Union[Assignment, FunctionCall, IfBlock, Break]
-TopLevelNode = Union[UseDirective, VarSection, ProcedureDecl, Section]
+ConditionExpr = Union[Comparison, LogicalOp, NotOp, Group, BoolLiteral, IdentCondition, FunctionCall]
+BodyNode = Union[Assignment, FunctionCall, IfBlock, Break, Comment]
+TopLevelNode = Union[UseDirective, VarSection, ProcedureDecl, Section, Comment]

tools/hrw4u/src/ast_visitor.py

@@ -29,6 +29,9 @@ class ASTVisitor(hrw4uVisitor):
     # method has an explicit return type and full control over how
     # child results are assembled into parent AST nodes.
 
+    def _visit_comment(self, ctx) -> Comment:
+        return Comment(text=ctx.COMMENT().getText(), line=ctx.start.line)
+
     def visitProgram(self, ctx) -> HRW4UAST:
         items = []
         for item in ctx.programItem():
@@ -39,7 +42,7 @@ def visitProgram(self, ctx) -> HRW4UAST:
             elif item.section() is not None:
                 items.append(self._visit_section(item.section()))
             elif item.commentLine() is not None:
-                pass
+                items.append(self._visit_comment(item.commentLine()))
             else:
                 raise ValueError(f"Unhandled programItem alternative at line {item.start.line}")
         return HRW4UAST(body=tuple(items))
@@ -75,7 +78,7 @@ def _visit_var_section(self, ctx, scope) -> VarSection:
             if var_item.variableDecl() is not None:
                 decls.append(self._visit_var_decl(var_item.variableDecl()))
             elif var_item.commentLine() is not None:
-                pass
+                decls.append(self._visit_comment(var_item.commentLine()))
             else:
                 raise ValueError(f"Unhandled variablesItem alternative at line {var_item.start.line}")
         return VarSection(scope=scope, declarations=tuple(decls), line=ctx.start.line)
@@ -93,7 +96,7 @@ def _visit_body(self, items) -> list[BodyNode]:
             elif item.conditional() is not None:
                 result.append(self._visit_conditional(item.conditional()))
             elif item.commentLine() is not None:
-                pass
+                result.append(self._visit_comment(item.commentLine()))
             else:
                 raise ValueError(f"Unhandled body item alternative at line {item.start.line}")
         return result
@@ -125,19 +128,19 @@ def _visit_function_call(self, ctx) -> FunctionCall:
 
     def _extract_value(self, ctx) -> ValueExpr:
         if ctx.number is not None:
-            return int(ctx.number.text)
+            return NumberValue(raw=ctx.number.text)
         if ctx.str_ is not None:
             return LiteralStringValue(raw=ctx.str_.text[1:-1])
         if ctx.TRUE():
-            return True
+            return BoolValue(raw=ctx.TRUE().getText())
         if ctx.FALSE():
-            return False
+            return BoolValue(raw=ctx.FALSE().getText())
         if ctx.ident is not None:
             return IdentValue(raw=ctx.ident.text)
         if ctx.ip():
             return IPValue(raw=ctx.ip().getText())
         if ctx.iprange():
-            return tuple(IPValue(raw=ip.getText()) for ip in ctx.iprange().ip())
+            return IpRangeValue(raw=ctx.iprange().getText())
         if ctx.paramRef():
             return ParamRef(raw=ctx.paramRef().IDENT().getText())
         raise ValueError(f"Unhandled value alternative at line {ctx.start.line}")
@@ -155,13 +158,17 @@ def _visit_conditional(self, ctx) -> IfBlock:
             elif_body = tuple(self._visit_body(elif_block.blockItem())) if elif_block else ()
             elif_branches.append(ElifBranch(condition=elif_cond, body=elif_body, line=elif_ctx.start.line))
 
-        else_body = ()
-        if ctx.elseClause():
-            else_block = ctx.elseClause().block()
-            if else_block:
-                else_body = tuple(self._visit_body(else_block.blockItem()))
+        else_clause = ctx.elseClause()
+        else_block = else_clause.block() if else_clause else None
+        else_body = tuple(self._visit_body(else_block.blockItem())) if else_block else ()
 
-        return IfBlock(condition=condition, body=body, elif_branches=tuple(elif_branches), else_body=else_body, line=ctx.start.line)
+        return IfBlock(
+            condition=condition,
+            body=body,
+            elif_branches=tuple(elif_branches),
+            else_body=else_body,
+            has_else=else_clause is not None,
+            line=ctx.start.line)
 
     def _visit_condition(self, ctx) -> ConditionExpr:
         return self._visit_expression(ctx.expression())
@@ -184,7 +191,7 @@ def _visit_factor(self, ctx) -> ConditionExpr:
         if ctx.getChildCount() == 2 and ctx.getChild(0).getText() == "!":
             return NotOp(operand=self._visit_factor(ctx.factor()), line=ctx.start.line)
         if ctx.LPAREN():
-            return self._visit_expression(ctx.expression())
+            return Group(inner=self._visit_expression(ctx.expression()), line=ctx.start.line)
         if ctx.functionCall():
             return self._visit_function_call(ctx.functionCall())
         if ctx.comparison():
@@ -231,14 +238,14 @@ def _detect_comparison_operator(self, ctx) -> str:
             return "in"
         raise ValueError(f"Unhandled comparison operator at line {ctx.start.line}")
 
-    def _extract_comparison_rhs(self, ctx, operator) -> ValueExpr | RegexValue | tuple[ValueExpr, ...]:
+    def _extract_comparison_rhs(self, ctx, operator) -> ValueExpr | RegexValue | SetValue | IpRangeValue:
         if operator in ("~", "!~"):
             return RegexValue(raw=ctx.regex().getText()[1:-1])
         if operator in ("in", "!in"):
             if ctx.set_():
-                return tuple(self._extract_value(v) for v in ctx.set_().value())
+                return SetValue(raw=ctx.set_().getText()[1:-1])
             if ctx.iprange():
-                return tuple(IPValue(raw=ip.getText()) for ip in ctx.iprange().ip())
+                return IpRangeValue(raw=ctx.iprange().getText())
         if ctx.value():
             return self._extract_value(ctx.value())
         raise ValueError(f"Unhandled comparison RHS at line {ctx.start.line}")

tools/hrw4u/tests/ast_unparse.py

@@ -0,0 +1,144 @@
+#
+#  Licensed to the Apache Software Foundation (ASF) under one
+#  or more contributor license agreements.  See the NOTICE file
+#  distributed with this work for additional information
+#  regarding copyright ownership.  The ASF licenses this file
+#  to you under the Apache License, Version 2.0 (the
+#  "License"); you may not use this file except in compliance
+#  with the License.  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+"""Render an AST back to hrw4u source. Test-only; see test_ast_roundtrip.py for why.
+
+Whitespace is not reproduced: the emitter derives its own indentation and never reads the
+source's.
+"""
+
+from __future__ import annotations
+
+from hrw4u.ast_nodes import *
+
+INDENT = "    "
+
+
+def unparse(ast: HRW4UAST) -> str:
+    return "\n".join(_top_level(item) for item in ast.body) + "\n"
+
+
+def _top_level(node: TopLevelNode) -> str:
+    match node:
+        case Comment():
+            return node.text
+        case UseDirective():
+            return f"use {node.spec}"
+        case ProcedureDecl():
+            params = ", ".join(_proc_param(p) for p in node.params)
+            return _braced(f"procedure {node.name}({params})", [_body_item(b, 1) for b in node.body])
+        case VarSection():
+            keyword = "SESSION_VARS" if node.scope == "session" else "VARS"
+            return _braced(keyword, [_var_item(d) for d in node.declarations])
+        case Section():
+            return _braced(node.type, [_body_item(b, 1) for b in node.body])
+    raise ValueError(f"unparse: unhandled top-level node {type(node).__name__}")
+
+
+def _braced(header: str, lines: list[str]) -> str:
+    return "\n".join([f"{header} {{", *(f"{INDENT}{line}" for line in lines), "}"])
+
+
+def _proc_param(p: ProcParam) -> str:
+    return f"${p.name}" if p.default is None else f"${p.name}={_value(p.default)}"
+
+
+def _var_item(node: VarDecl | Comment) -> str:
+    if isinstance(node, Comment):
+        return node.text
+    slot = "" if node.slot is None else f" @{node.slot}"
+    return f"{node.name}: {node.type_name}{slot};"
+
+
+def _body_item(node: BodyNode, depth: int) -> str:
+    match node:
+        case Comment():
+            return node.text
+        case Break():
+            return "break;"
+        case FunctionCall():
+            return f"{_call(node)};"
+        case Assignment():
+            return f"{_target(node.target)} {node.operator} {_value(node.value)};"
+        case IfBlock():
+            return _if_block(node, depth)
+    raise ValueError(f"unparse: unhandled body node {type(node).__name__}")
+
+
+def _if_block(node: IfBlock, depth: int) -> str:
+    pad = INDENT * depth
+    lines = [f"if {_condition(node.condition)} {{"]
+    lines += [f"{INDENT}{line}" for line in _nested(node.body, depth)]
+    for arm in node.elif_branches:
+        lines.append(f"}} elif {_condition(arm.condition)} {{")
+        lines += [f"{INDENT}{line}" for line in _nested(arm.body, depth)]
+    if node.has_else:
+        lines.append("} else {")
+        lines += [f"{INDENT}{line}" for line in _nested(node.else_body, depth)]
+    lines.append("}")
+    return f"\n{pad}".join(lines)
+
+
+def _nested(body: tuple[BodyNode, ...], depth: int) -> list[str]:
+    return [line for item in body for line in _body_item(item, depth + 1).splitlines()]
+
+
+def _target(t: Target) -> str:
+    return t.field if t.namespace is None else f"{t.namespace}.{t.field}"
+
+
+def _call(node: FunctionCall) -> str:
+    return f"{node.name}({', '.join(_value(a) for a in node.args)})"
+
+
+def _condition(node: ConditionExpr) -> str:
+    match node:
+        case Group():
+            return f"({_condition(node.inner)})"
+        case LogicalOp():
+            return f"{_condition(node.left)} {node.operator} {_condition(node.right)}"
+        case NotOp():
+            return f"!{_condition(node.operand)}"
+        case BoolLiteral():
+            return "true" if node.value else "false"
+        case IdentCondition():
+            return node.name
+        case FunctionCall():
+            return _call(node)
+        case Comparison():
+            return _comparison(node)
+    raise ValueError(f"unparse: unhandled condition node {type(node).__name__}")
+
+
+def _comparison(node: Comparison) -> str:
+    left = node.left.raw if isinstance(node.left, IdentValue) else _call(node.left)
+    mods = f" with {', '.join(node.modifiers)}" if node.modifiers else ""
+    return f"{left} {node.operator} {_value(node.right)}{mods}"
+
+
+def _value(v: ValueExpr | RegexValue | SetValue) -> str:
+    match v:
+        case LiteralStringValue():
+            return f'"{v.raw}"'
+        case NumberValue() | BoolValue() | IdentValue() | IPValue() | IpRangeValue():
+            return v.raw
+        case ParamRef():
+            return f"${v.raw}"
+        case RegexValue():
+            return f"/{v.raw}/"
+        case SetValue():
+            return f"[{v.raw}]"
+    raise ValueError(f"unparse: unhandled value {type(v).__name__}")

tools/hrw4u/tests/data/ops/bool-spelling.input.txt

@@ -0,0 +1,7 @@
+# The emitter echoes a bool's spelling, so the AST cannot normalize one anywhere:
+# an assignment RHS is not the only value context that reaches header_rewrite.
+REMAP {
+    if inbound.req.X-Debug == TRUE {
+        set-config("proxy.config.http.cache.http", FALSE);
+    }
+}

tools/hrw4u/tests/data/ops/bool-spelling.output.txt

@@ -0,0 +1,5 @@
+# The emitter echoes a bool's spelling, so the AST cannot normalize one anywhere:
+# an assignment RHS is not the only value context that reaches header_rewrite.
+cond %{REMAP_PSEUDO_HOOK} [AND]
+cond %{CLIENT-HEADER:X-Debug} =TRUE
+    set-config "proxy.config.http.cache.http" FALSE

tools/hrw4u/tests/data/ops/exceptions.txt

@@ -9,3 +9,5 @@ header_value_context.input: u4wrh
 http_cntl_valid_bools.input: hrw4u
 # Literal JSON blocks compile to native header_rewrite escaping.
 json-body.input: hrw4u
+# The reverse normalizes a bool argument's spelling (FALSE -> false)
+bool-spelling.input: hrw4u

tools/hrw4u/tests/data/ops/number-spelling.input.txt

@@ -0,0 +1,7 @@
+# The emitter echoes a number's digits, so the AST cannot normalize one anywhere:
+# 007 and 7 are different bytes once a number reaches a header value.
+REMAP {
+    if random(0100) > 007 {
+        inbound.req.X-Count = 007;
+    }
+}