Add $ref support to policy schema

components/policy/tools/generate_policy_source.py

Index: components/policy/tools/generate_policy_source.py
diff --git a/components/policy/tools/generate_policy_source.py b/components/policy/tools/generate_policy_source.py
index 9ccf881d17137a2ecd5b98f6e3ea77f21a44cda8..0be012abf272e7c6998a84d851e007416443d5ec 100755
--- a/components/policy/tools/generate_policy_source.py
+++ b/components/policy/tools/generate_policy_source.py
@@ -11,11 +11,13 @@ chromium_os_flag should be 1 if this is a Chromium OS build
 template is the path to a .json policy template file.'''
 
 from __future__ import with_statement
+from functools import partial
 import json
 from optparse import OptionParser
 import re
 import sys
 import textwrap
+import types
 
 
 CHROME_POLICY_KEY = 'SOFTWARE\\\\Policies\\\\Google\\\\Chrome'
@@ -282,6 +284,7 @@ class SchemaNodesGenerator:
     }
     self.stringlist_type = None
     self.ranges = {}
+    self.id_map = {}
 
   def GetString(self, s):
     if s in self.shared_strings:
@@ -398,6 +401,12 @@ class SchemaNodesGenerator:
 
     |schema|: a valid JSON schema in a dictionary.
     |name|: the name of the current node, for the generated comments."""
+    if schema.has_key('$ref'):
+      if schema.has_key('id'):
+        raise RuntimeError("Schemas with a $ref can't have an id")
+      if not isinstance(schema['$ref'], types.StringTypes):
+        raise RuntimeError("$ref attribute must be a string")
+      return schema['$ref']
     if schema['type'] in self.simple_types:
       if not self.SchemaHaveRestriction(schema):
         # Simple types use shared nodes.
@@ -413,9 +422,8 @@ class SchemaNodesGenerator:
       # Special case for lists of strings, which is a common policy type.
       if schema['items']['type'] == 'string':
         return self.GetStringList()
-      return self.AppendSchema(
-          'TYPE_LIST',
-          self.Generate(schema['items'], 'items of ' + name))
+      return self.AppendSchema('TYPE_LIST',
+          self.GenerateAndCollectID(schema['items'], 'items of ' + name))
     elif schema['type'] == 'object':
       # Reserve an index first, so that dictionaries come before their
       # properties. This makes sure that the root node is the first in the
@@ -423,7 +431,7 @@ class SchemaNodesGenerator:
       index = self.AppendSchema('TYPE_DICTIONARY', -1)
 
       if 'additionalProperties' in schema:
-        additionalProperties = self.Generate(
+        additionalProperties = self.GenerateAndCollectID(
             schema['additionalProperties'],
             'additionalProperties of ' + name)
       else:
@@ -434,13 +442,14 @@ class SchemaNodesGenerator:
       # recursive calls to Generate() append the necessary child nodes; if
       # |properties| were a generator then this wouldn't work.
       sorted_properties = sorted(schema.get('properties', {}).items())
-      properties = [ (self.GetString(key), self.Generate(subschema, key))
-                     for key, subschema in sorted_properties ]
+      properties = [
+          (self.GetString(key), self.GenerateAndCollectID(subschema, key))
+          for key, subschema in sorted_properties ]
 
       pattern_properties = []
       for pattern, subschema in schema.get('patternProperties', {}).items():
         pattern_properties.append((self.GetString(pattern),
-            self.Generate(subschema, pattern)));
+            self.GenerateAndCollectID(subschema, pattern)));
 
       begin = len(self.property_nodes)
       self.property_nodes += properties
@@ -462,6 +471,20 @@ class SchemaNodesGenerator:
     else:
       assert False
 
+  def GenerateAndCollectID(self, schema, name):
+    """A wrapper of Generate(), will take the return value, check and add 'id'
+    attribute to self.id_map. The wrapper needs to be used for every call to
+    Generate().
+    """
+    index = self.Generate(schema, name)
+    if not schema.has_key('id'):
+      return index
+    id_str = schema['id']
+    if self.id_map.has_key(id_str):
+      raise RuntimeError('Duplicated id: ' + id_str)
+    self.id_map[id_str] = index
+    return index
+
   def Write(self, f):
     """Writes the generated structs to the given file.
 
@@ -516,6 +539,29 @@ class SchemaNodesGenerator:
     f.write('  kStringEnumerations,\n' if self.string_enums else '  NULL,\n')
     f.write('};\n\n')
 
+  def GetByID(self, id_str):
+    if not isinstance(id_str, types.StringTypes):
+      return id_str
+    if not self.id_map.has_key(id_str):
+      raise RuntimeError('Invalid $ref: ' + id_str)
+    return self.id_map[id_str]
+
+  def ResolveID(self, index, params):
+    return params[:index] + (self.GetByID(params[index]),) + params[index+1:]
+
+  def ResolveReferences(self):
+    """Resolve reference mapping, required to be called after Generate()
+
+    After calling Generate(), the type of indices used in schema structures
+    might be either int or string. An int type suggests that it's a resolved
+    index, but for string type it's unresolved. Resolving a reference is as
+    simple as looking up for corresponding ID in self.id_map, and replace the
+    old index with the mapped index.
+    """
+    self.schema_nodes = map(partial(self.ResolveID, 1), self.schema_nodes)
+    self.property_nodes = map(partial(self.ResolveID, 1), self.property_nodes)
+    self.properties_nodes = map(partial(self.ResolveID, 3),
+        self.properties_nodes)
 
 def _WritePolicyConstantSource(policies, os, f):
   f.write('#include "policy/policy_constants.h"\n'
@@ -558,7 +604,8 @@ def _WritePolicyConstantSource(policies, os, f):
   f.write('};\n\n')
 
   schema_generator = SchemaNodesGenerator(shared_strings)
-  schema_generator.Generate(chrome_schema, 'root node')
+  schema_generator.GenerateAndCollectID(chrome_schema, 'root node')
+  schema_generator.ResolveReferences()
   schema_generator.Write(f)
 
   f.write('bool CompareKeys(const internal::PropertyNode& node,\n'