chore: import upstream snapshot with attribution

python/paddle/distributed/flex_checkpoint/aoa/parser.py

@@ -0,0 +1,142 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed 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.
+
+from .lexer import Token, TokenType
+
+
+class Statement:
+    def __init__(self, left_vars, right_vars, attrs):
+        self.left_vars = left_vars  # List[Var]
+        self.right_vars = right_vars  # List[Var]
+        self.attrs = attrs  # List[Attribute]
+
+    def __repr__(self):
+        return f"Statement({self.left_vars} -> {self.right_vars}, attrs={self.attrs})"
+
+
+class Var:
+    def __init__(self, name):
+        self.name = name
+
+    def __repr__(self):
+        return self.name
+
+
+class Attribute:
+    def __init__(self, key, value):
+        self.key = key
+        self.value = value
+
+    def __repr__(self):
+        return f"{self.key}={self.value!r}"
+
+
+class Parser:
+    """
+    AOA Grammar
+    PROGRAM   ::= { STATEMENT }
+
+    STATEMENT ::= VAR_LIST '->' VAR ',' ATTR_LIST       // meige
+                | VAR '->' VAR_LIST ',' ATTR_LIST       // split
+                | VAR '->' VAR ',' ATTR_LIST            // single variable mapping + attributes
+                | VAR '->' VAR                          // single variable mapping, rename
+
+    VAR_LIST  ::= VAR { ',' VAR }
+    VAR       ::= IDENTIFIER
+    ATTR_LIST ::= ATTRIBUTE { ',' ATTRIBUTE }
+    ATTRIBUTE ::= IDENTIFIER '=' VALUE
+    VALUE     ::= NUMBER | STRING
+    """
+
+    def __init__(self, tokens):
+        self.tokens = tokens
+        self.pos = 0
+
+    def at_end(self):
+        return self.peek().type == TokenType.EOF
+
+    def peek(self, offset=0):
+        if self.pos + offset >= len(self.tokens):
+            return Token(TokenType.EOF, '')
+        return self.tokens[self.pos + offset]
+
+    def consume(self, expected_type=None):
+        tok = self.peek()
+        if expected_type and tok.type != expected_type:
+            raise SyntaxError(
+                f'Expected {expected_type}, got {tok.type} at pos {self.pos}'
+            )
+        self.pos += 1
+        return tok
+
+    def expect(self, expected_type):
+        return self.consume(expected_type)
+
+    def skip_newlines(self):
+        while self.peek().type == TokenType.NEWLINE:
+            self.consume()
+
+    def parse_program(self):
+        stmts = []
+        self.skip_newlines()
+        while not self.at_end():
+            stmt = self.parse_statement()
+            stmts.append(stmt)
+            self.skip_newlines()
+        return stmts
+
+    def parse_statement(self):
+        left_vars = [self.parse_var()]
+        while self.peek().type == TokenType.COMMA:
+            self.consume(TokenType.COMMA)
+            left_vars.append(self.parse_var())
+        self.expect(TokenType.RARROW)
+        right_vars = [self.parse_var()]
+        while self.peek().type == TokenType.COMMA:
+            # Lookahead for attribute: IDENT '=' after COMMA means attribute starts
+            if (
+                self.peek(1).type == TokenType.IDENTIFIER
+                and self.peek(2).type == TokenType.EQUAL
+            ):
+                break
+            self.consume(TokenType.COMMA)
+            right_vars.append(self.parse_var())
+        attrs = []
+        if self.peek().type == TokenType.COMMA:
+            self.consume(TokenType.COMMA)
+            attrs = self.parse_attr_list()
+        return Statement(left_vars, right_vars, attrs)
+
+    def parse_var(self):
+        name = self.expect(TokenType.IDENTIFIER).value
+        return Var(name)
+
+    def parse_attr_list(self):
+        attrs = [self.parse_attribute()]
+        while self.peek().type == TokenType.COMMA:
+            self.consume(TokenType.COMMA)
+            attrs.append(self.parse_attribute())
+        return attrs
+
+    def parse_attribute(self):
+        key = self.expect(TokenType.IDENTIFIER).value
+        self.expect(TokenType.EQUAL)
+        val_tok = self.consume()
+        if val_tok.type == TokenType.NUMBER:
+            val = int(val_tok.value)
+        elif val_tok.type == TokenType.STRING:
+            val = val_tok.value.strip('"').strip("'")
+        else:
+            raise SyntaxError(f'Unexpected value: {val_tok}')
+        return Attribute(key, val)