TransPyC/lib/core/Translator/ConstEval.py

from __future__ import annotations

import ast
import llvmlite.ir as ir


class ConstEvalMixin:
    """常量表达式求值 Mixin

    提供数组初始化常量构建与常量表达式求值能力。
    """

    def _BuildArrayInitConstants(self, value_node, ElemType, elem_type_node, ArrayCount, Gen):
        from lib.core.Handles.HandlesAssign import AssignHandle
        _zv = AssignHandle._zero_value(ElemType)
        _fallback = _zv if _zv is not None else ir.Constant(ElemType, ir.Undefined)
        if not isinstance(value_node, (ast.List, ast.Set)):
            return None
        if isinstance(ElemType, ir.ArrayType):
            constants = []
            for elt in value_node.elts:
                if isinstance(elt, (ast.List, ast.Set)):
                    inner_consts = self._BuildMultiDimArrayInitConstants(elt, ElemType, Gen)
                    if inner_consts:
                        try:
                            constants.append(ir.Constant(ElemType, inner_consts))
                        except Exception:  # fallback to _fallback on constant construction failure
                            constants.append(_fallback)
                    else:
                        constants.append(_fallback)
                else:
                    constants.append(_fallback)
            while len(constants) < ArrayCount:
                constants.append(_fallback)
            return constants[:ArrayCount]
        StructName = None
        if isinstance(elem_type_node, ast.Name) and elem_type_node.id in Gen.structs:
            StructName = elem_type_node.id
        constants = []
        for elt in value_node.elts:
            if StructName and isinstance(elt, ast.Call) and isinstance(elt.func, ast.Name) and elt.func.id == StructName:
                const = self.ClassHandler._BuildStructConstant(elt, StructName, Gen)
                if const:
                    constants.append(const)
                else:
                    constants.append(_fallback)
            elif isinstance(elt, ast.Constant):
                const = self.ClassHandler._BuildScalarConstant(elt, ElemType)
                if const:
                    constants.append(const)
                else:
                    constants.append(_fallback)
            elif isinstance(elt, ast.UnaryOp):
                const = self.ClassHandler._BuildScalarConstant(elt, ElemType)
                if const:
                    constants.append(const)
                else:
                    constants.append(_fallback)
            elif isinstance(elt, ast.Call) and isinstance(elt.func, ast.Name) and elt.func.id == 'chr':
                if elt.args and isinstance(elt.args[0], ast.Constant) and isinstance(elt.args[0].value, int):
                    try:
                        constants.append(ir.Constant(ElemType, elt.args[0].value & 0xFF))
                    except Exception:  # fallback to _fallback on constant construction failure
                        constants.append(_fallback)
                else:
                    constants.append(_fallback)
            elif isinstance(elt, ast.Call) and isinstance(elt.func, ast.Name) and elt.func.id == 'ord':
                if elt.args and isinstance(elt.args[0], ast.Constant) and isinstance(elt.args[0].value, str) and len(elt.args[0].value) == 1:
                    try:
                        constants.append(ir.Constant(ElemType, ord(elt.args[0].value)))
                    except Exception:  # fallback to _fallback on constant construction failure
                        constants.append(_fallback)
                else:
                    constants.append(_fallback)
            else:
                const_val = self._try_eval_const_call(elt, Gen)
                if const_val is not None:
                    try:
                        constants.append(ir.Constant(ElemType, const_val))
                    except Exception:  # fallback to _fallback on constant construction failure
                        constants.append(_fallback)
                else:
                    constants.append(_fallback)
        while len(constants) < ArrayCount:
            constants.append(_fallback)
        return constants[:ArrayCount]

    def _try_eval_const_call(self, node, Gen):
        import ast
        if isinstance(node, ast.Call):
            func_name = None
            if isinstance(node.func, ast.Name):
                func_name = node.func.id
            elif isinstance(node.func, ast.Attribute):
                func_name = node.func.attr
            if func_name:
                arg_vals = []
                for arg in node.args:
                    v = self._try_eval_const_expr(arg, Gen)
                    if v is None:
                        return None
                    arg_vals.append(v)
                if func_name == 'COLOR_RGB' and len(arg_vals) == 3:
                    r, g, b = arg_vals
                    return (255 << 24) | (int(b) << 16) | (int(g) << 8) | int(r)
                if func_name in Gen._DefineConstants:
                    return Gen._DefineConstants[func_name]
                for key in Gen.functions:
                    if key.endswith(f'__{func_name}'):
                        return None
            return None
        return self._try_eval_const_expr(node, Gen)

    def _try_eval_const_expr(self, node, Gen):
        from lib.core.ConstEvaluator import ConstEvaluator
        return ConstEvaluator.eval_with_symtab(node, self.SymbolTable)

    def _BuildMultiDimArrayInitConstants(self, value_node, ArrayType, Gen):
        from lib.core.Handles.HandlesAssign import AssignHandle
        inner_type = ArrayType.element
        count = ArrayType.count
        elts = list(value_node.elts) if hasattr(value_node, 'elts') else []
        _zv = AssignHandle._zero_value(inner_type)
        _fallback = _zv if _zv is not None else ir.Constant(inner_type, ir.Undefined)
        constants = []
        for elt in elts:
            if isinstance(inner_type, ir.ArrayType):
                if isinstance(elt, (ast.List, ast.Set)):
                    inner_consts = self._BuildMultiDimArrayInitConstants(elt, inner_type, Gen)
                    if inner_consts:
                        constants.append(ir.Constant(inner_type, inner_consts))
                    else:
                        constants.append(_fallback)
                else:
                    constants.append(_fallback)
            elif isinstance(elt, ast.Constant):
                const = self.ClassHandler._BuildScalarConstant(elt, inner_type)
                if const:
                    constants.append(const)
                else:
                    constants.append(_fallback)
            elif isinstance(elt, ast.UnaryOp):
                const = self.ClassHandler._BuildScalarConstant(elt, inner_type)
                if const:
                    constants.append(const)
                else:
                    constants.append(_fallback)
            else:
                constants.append(_fallback)
        while len(constants) < count:
            constants.append(_fallback)
        return constants[:count]