lunaix-os/scripts/build-tools/lcfg2/rewriter.py

   1 import ast
   2
   3 from lib.utils  import Schema
   4 from .lazy       import Lazy
   5 from .common     import NodeProperty
   6
   7 class RewriteRule:
   8     MaybeBuiltin = Schema(
   9                         ast.Call,
  10                             func=Schema(ast.Name),
  11                             args=[ast.expr])
  12
  13     WhenTogglerItem = Schema(
  14                             ast.Compare,
  15                                 left=ast.Name,
  16                                 ops=[Schema.Union(ast.Is, ast.IsNot)],
  17                                 comparators=[ast.Constant])
  18
  19     WhenToggler = Schema(
  20                     Schema.Union(
  21                         WhenTogglerItem,
  22                         Schema(ast.BoolOp,
  23                             op=ast.And,
  24                             values=Schema.List(WhenTogglerItem))))
  25
  26 class ConfigNodeASTRewriter(ast.NodeTransformer):
  27
  28
  29     def __init__(self, cfg_node):
  30         super().__init__()
  31
  32         self.__cfg_node = cfg_node
  33
  34     def __subscript_accessor(self, name, ctx, token):
  35         return ast.Subscript(
  36             value=ast.Name("__lzLut__", ctx=ast.Load()),
  37             slice=ast.Constant(name),
  38             ctx=ctx,
  39         )
  40
  41     def __gen_accessor(self, orig):
  42         key = Lazy.from_astn(self.__cfg_node, orig)
  43         if not key:
  44             return self.generic_visit(orig)
  45
  46         return self.__subscript_accessor(key, orig.ctx, orig)
  47
  48     def __gen_dependency(self, node):
  49         cfgn = self.__cfg_node
  50         dep_expr = NodeProperty.Dependency[cfgn]
  51         if dep_expr is None:
  52             NodeProperty.Dependency[cfgn] = node
  53             return
  54
  55         if not isinstance(dep_expr, ast.expr):
  56             raise cfgn.config_error(
  57                 f"invalid dependency state: {dep_expr}")
  58
  59         dep_expr = ast.BoolOp(ast.And(), [dep_expr, node])
  60         NodeProperty.Dependency[cfgn] = dep_expr
  61
  62     def __gen_when_expr(self, node):
  63         and_list = []
  64         cfgn = self.__cfg_node
  65
  66         if RewriteRule.WhenToggler != node:
  67             raise cfgn.config_error(
  68                 f"invalid when(...) expression: {ast.unparse(node)}")
  69
  70         if RewriteRule.WhenTogglerItem == node:
  71             and_list.append(node)
  72         else:
  73             and_list += node.values
  74
  75         for i in range(len(and_list)):
  76             item = and_list[i]
  77             operator = item.ops[0]
  78
  79             name = Lazy.from_type(cfgn, Lazy.NodeValue, item.left.id)
  80             acc = self.__subscript_accessor(name, ast.Load(), node)
  81
  82             if isinstance(operator, ast.Is):
  83                 operator = ast.Eq()
  84             else:
  85                 operator = ast.NotEq()
  86
  87             item.left = acc
  88             item.ops  = [operator]
  89
  90
  91         current = ast.BoolOp(
  92                         op=ast.And(),
  93                         values=[ast.Constant(True), *and_list])
  94
  95         expr = NodeProperty.WhenToggle[cfgn]
  96         if expr:
  97             assert isinstance(expr, ast.expr)
  98             current = ast.BoolOp(op=ast.Or(), values=[expr, current])
  99
 100         NodeProperty.WhenToggle[cfgn] = current
 101
 102     def visit_Attribute(self, node):
 103         return self.__gen_accessor(node)
 104
 105     def visit_Expr(self, node):
 106         val = node.value
 107
 108         if RewriteRule.MaybeBuiltin != val:
 109             return self.generic_visit(node)
 110
 111         # Process marker functions
 112         name = val.func.id
 113         if name == "require":
 114             self.__gen_dependency(val.args[0])
 115         elif name == "when":
 116             self.__gen_when_expr(val.args[0])
 117         else:
 118             return self.generic_visit(node)
 119
 120         return None
 121
 122     def visit_Return(self, node):
 123         if NodeProperty.WhenToggle[self.__cfg_node]:
 124             return None
 125         return self.generic_visit(node)
 126
 127     def visit_Is(self, node):
 128         return ast.Eq()
 129
 130     def rewrite(self, node):
 131         assert isinstance(node, ast.Module)
 132         node = self.visit(node)
 133
 134         expr = NodeProperty.WhenToggle[self.__cfg_node]
 135         if not expr:
 136             return node
 137
 138         del NodeProperty.WhenToggle[self.__cfg_node]
 139
 140         node.body.append(ast.Return(expr, lineno=0, col_offset=0))
 141         return node