Refactor geckolib atoms regen script

1 files changed, 67 insertions, 43 deletions

diff --git a/components/style/binding_tools/regen_atoms.py b/components/style/binding_tools/regen_atoms.py
index 19b72799edf..1745fe8a84d 100755
--- a/components/style/binding_tools/regen_atoms.py
+++ b/components/style/binding_tools/regen_atoms.py
@@ -8,18 +8,27 @@ import re
 import os
 
 
+PRELUDE = """
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+/* Autogenerated file, DO NOT EDIT DIRECTLY */
+"""[1:]
+
+
 def gnu_symbolify(source, ident):
-    return "_ZN" + str(len(source.CLASS)) + source.CLASS + str(len(ident)) + ident + "E"
+    return "_ZN{}{}{}{}E".format(len(source.CLASS), source.CLASS, len(ident), ident)
 
 
 def msvc64_symbolify(source, ident):
-    return "?" + ident + "@" + source.CLASS + "@@2PEAV" + source.TYPE + "@@EA"
+    return "?{}@{}@@2PEAV{}@@EA".format(ident, source.CLASS, source.TYPE)
 
 
 def msvc32_symbolify(source, ident):
     # Prepend "\x01" to avoid LLVM prefixing the mangled name with "_".
     # See https://github.com/rust-lang/rust/issues/36097
-    return "\\x01?" + ident + "@" + source.CLASS + "@@2PAV" + source.TYPE + "@@A"
+    return "\\x01?{}@{}@@2PAV{}@@A".format(ident, source.CLASS, source.TYPE)
 
 
 class GkAtomSource:
@@ -92,56 +101,71 @@ def collect_atoms(objdir):
                     atoms.append(Atom(source, result.group(1), result.group(2)))
     return atoms
 
-PRELUDE = """
-/* This Source Code Form is subject to the terms of the Mozilla Public
- * License, v. 2.0. If a copy of the MPL was not distributed with this
- * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
 
-/* Autogenerated file, DO NOT EDIT DIRECTLY */
-"""[1:]
+IMPORTS = ("\nuse gecko_bindings::structs::nsIAtom;"
+           "\nuse string_cache::Atom;\n\n")
+
+ATOM_TEMPLATE = ("            #[link_name = \"{link_name}\"]\n"
+                 "            pub static {name}: *mut {type};")
+
+UNSAFE_STATIC = ("#[inline(always)]\n"
+                 "pub fn unsafe_atom_from_static(ptr: *mut nsIAtom) -> Atom {\n"
+                 "    unsafe { Atom::from_static(ptr) }\n"
+                 "}\n\n")
+
+CFG_IF = '''
+cfg_if! {{
+    if #[cfg(not(target_env = "msvc"))] {{
+        extern {{
+{gnu}
+        }}
+    }} else if #[cfg(target_pointer_width = "64")] {{
+        extern {{
+{msvc64}
+        }}
+    }} else {{
+        extern {{
+{msvc32}
+        }}
+    }}
+}}
+'''
+
+RULE_TEMPLATE = ('("{atom}") => '
+                 '{{ $crate::string_cache::atom_macro::unsafe_atom_from_static'
+                 '($crate::string_cache::atom_macro::{name} as *mut _) }};')
+
+MACRO = '''
+#[macro_export]
+macro_rules! atom {{
+{}
+}}
+'''
 
 
 def write_atom_macro(atoms, file_name):
-    ATOM_TEMPLATE = """
-            #[link_name = "{link_name}"]
-            pub static {name}: *mut {type};
-"""[1:]
-
-    def write_items(f, func):
-        f.write("        extern {\n")
-        for atom in atoms:
-            f.write(ATOM_TEMPLATE.format(name=atom.ident,
-                                         link_name=func(atom),
-                                         type=atom.type()))
-        f.write("        }\n")
+    def get_symbols(func):
+        return '\n'.join([ATOM_TEMPLATE.format(name=atom.ident,
+                                               link_name=func(atom),
+                                               type=atom.type()) for atom in atoms])
 
     with open(file_name, "wb") as f:
         f.write(PRELUDE)
-        f.write("use gecko_bindings::structs::nsIAtom;\n\n")
-        f.write("use string_cache::Atom;\n\n")
+        f.write(IMPORTS)
+
         for source in SOURCES:
             if source.TYPE != "nsIAtom":
                 f.write("pub enum {} {{}}\n\n".format(source.TYPE))
-        f.write("""
-            #[inline(always)] pub fn unsafe_atom_from_static(ptr: *mut nsIAtom) -> Atom {
-                unsafe { Atom::from_static(ptr) }
-            }\n\n
-            """)
-        f.write("cfg_if! {\n")
-        f.write("    if #[cfg(not(target_env = \"msvc\"))] {\n")
-        write_items(f, Atom.gnu_symbol)
-        f.write("    } else if #[cfg(target_pointer_width = \"64\")] {\n")
-        write_items(f, Atom.msvc64_symbol)
-        f.write("    } else {\n")
-        write_items(f, Atom.msvc32_symbol)
-        f.write("    }\n")
-        f.write("}\n\n")
-        f.write("#[macro_export]\n")
-        f.write("macro_rules! atom {\n")
-        f.writelines(['("%s") => { $crate::string_cache::atom_macro::unsafe_atom_from_static(\
-                      $crate::string_cache::atom_macro::%s as *mut _) };\n'
-                     % (atom.value, atom.ident) for atom in atoms])
-        f.write("}\n")
+
+        f.write(UNSAFE_STATIC)
+
+        gnu_symbols = get_symbols(Atom.gnu_symbol)
+        msvc32_symbols = get_symbols(Atom.msvc32_symbol)
+        msvc64_symbols = get_symbols(Atom.msvc64_symbol)
+        f.write(CFG_IF.format(gnu=gnu_symbols, msvc32=msvc32_symbols, msvc64=msvc64_symbols))
+
+        macro_rules = [RULE_TEMPLATE.format(atom=atom.value, name=atom.ident) for atom in atoms]
+        f.write(MACRO.format('\n'.join(macro_rules)))
 
 
 PSEUDO_ELEMENT_HEADER = """