Create a compiled type for categories and use it in the DSL

TN: T116-036

langkit/compiled_types.py

@@ -3853,6 +3853,8 @@ def create_builtin_types():
              value_names=[names.Name('Unit_Specification'),
                           names.Name('Unit_Body')])
 
+    CompiledType('RefCategories', null_allowed=False)
+
     EnumType(name='LookupKind',
              location=None,
              doc="""

langkit/dsl_unparse.py

@@ -161,7 +161,7 @@ def emit_expr(expr, **ctx):
         Quantifier, If, IsNull, Cast, DynamicVariable, IsA, Not, SymbolLiteral,
         No, Cond, New, CollectionSingleton, Concat, EnumLiteral, EnvGet,
         ArrayLiteral, Arithmetic, PropertyError, CharacterLiteral,
-        StructUpdate, BigIntLiteral
+        StructUpdate, BigIntLiteral, RefCategories
     )
 
     then_underscore_var = ctx.get('then_underscore_var')
@@ -364,6 +364,8 @@ def emit_expr(expr, **ctx):
             args.append("from={}".format(ee(expr.sequential_from)))
         if expr.only_first:
             args.append("only_first={}".format(ee(expr.only_first)))
+        if expr.categories:
+            args.append('categories={}'.format(ee(expr.categories)))
         return emit_method_call(ee(expr.env), "get", args)
 
     elif is_a("bind"):
@@ -432,6 +434,11 @@ def emit_expr(expr, **ctx):
         return 'BigInt({})'.format(str(expr.expr)
                                    if isinstance(expr.expr, (int, long)) else
                                    ee(expr.expr))
+    elif isinstance(expr, RefCategories):
+        return 'RefCats({}, others={})'.format(', '.join(
+            '{}={}'.format(name, ee(value))
+            for name, value in sorted(expr.cat_map.items())
+        ), ee(expr.default))
     else:
         # raise NotImplementedError(type(expr))
         return repr(expr)

langkit/expressions/envs.py

@@ -4,13 +4,70 @@ from langkit import names
 from langkit.compiled_types import T, get_context
 from langkit.diagnostics import check_source_language
 from langkit.expressions.base import (
-    AbstractExpression, AbstractVariable, CallExpr, ComputingExpr,
-    GetSymbol, No, NullCheckExpr, NullExpr,
-    PropertyDef, Self, attr_call, auto_attr, construct
+    AbstractExpression, AbstractVariable, BindableLiteralExpr, CallExpr,
+    ComputingExpr, GetSymbol, No, NullCheckExpr, NullExpr, PropertyDef, Self,
+    attr_call, auto_attr, construct, dsl_document
 )
 from langkit.expressions.utils import assign_var
 
 
+@dsl_document
+class RefCategories(AbstractExpression):
+    """
+    Build a set of categories.
+    """
+
+    class Expr(BindableLiteralExpr):
+        def __init__(self, cats, abstract_expr=None):
+            self.cats = cats
+            super(RefCategories.Expr, self).__init__(
+                self.render_private_ada_constant(),
+                T.RefCategories,
+                abstract_expr=abstract_expr
+            )
+
+        def render_private_ada_constant(self):
+            all_cats = get_context().ref_cats
+            return '({})'.format(', '.join(sorted(
+                '{} => {}'.format(name.camel_with_underscores,
+                                  name in self.cats)
+                for name in all_cats
+            )))
+
+        # This type is not available in public APIs, so there is no need to
+        # implement the other rendering properties.
+
+        @property
+        def subexprs(self):
+            return {'cats': self.cats}
+
+    def __init__(self, default=False, **kwargs):
+        super(RefCategories, self).__init__()
+        self.default = default
+        self.cat_map = kwargs
+
+    def construct(self):
+        check_source_language(isinstance(self.default, bool),
+                              'Invalid categories default')
+
+        all_cats = get_context().ref_cats
+        cats = set(all_cats) if self.default else set()
+
+        # Compute the list of requested categories
+        for key, value in self.cat_map.items():
+            name = names.Name.from_lower(key)
+            check_source_language(name in all_cats,
+                                  'Invalid category: {}'.format(key))
+            check_source_language(isinstance(value, bool),
+                                  'Invalid status for {}'.format(key))
+            if value:
+                cats.add(name)
+            else:
+                cats.discard(name)
+
+        return self.Expr(cats, abstract_expr=self)
+
+
 @attr_call('get')
 def get(env, symbol, lookup=None, from_node=None, categories=None):
     """
@@ -58,10 +115,9 @@ class EnvGet(AbstractExpression):
     """
 
     class Expr(ComputingExpr):
-        def __init__(self, env_expr, key_expr, lookup_kind_expr,
+        def __init__(self, env_expr, key_expr, lookup_kind_expr, categories,
                      sequential_from=None,
-                     only_first=False, abstract_expr=None,
-                     categories=None):
+                     only_first=False, abstract_expr=None):
             self.env_expr = env_expr
             self.key_expr = key_expr
             self.lookup_kind_expr = lookup_kind_expr
@@ -80,25 +136,18 @@ class EnvGet(AbstractExpression):
             PropertyDef.get().set_uses_envs()
 
         def _render_pre(self):
-            if self.categories:
-                cat_arg = "({})".format(", ".join(
-                    "{} => {}".format(cat_name, cat_val)
-                    for cat_name, cat_val in self.categories.items()
-                ))
-            else:
-                cat_arg = "All_Cats"
-
             result = [
                 self.env_expr.render_pre(),
                 self.key_expr.render_pre(),
                 self.lookup_kind_expr.render_pre(),
+                self.categories.render_pre()
             ]
             args = [('Self', self.env_expr.render_expr()),
                     ('Key', self.key_expr.render_expr()),
                     ('Lookup_Kind', 'To_Lookup_Kind_Type ({})'.format(
                         self.lookup_kind_expr.render_expr()
                     )),
-                    ('Categories', cat_arg)]
+                    ('Categories', self.categories.render_expr())]
 
             # Pass the From parameter if the user wants sequential semantics
             if self.sequential_from:
@@ -131,6 +180,7 @@ class EnvGet(AbstractExpression):
                 'env': self.env_expr,
                 'key': self.key_expr,
                 'lookup_kind': self.lookup_kind_expr,
+                'categories': self.categories,
                 'sequential_from': self.sequential_from,
             }
 
@@ -188,38 +238,19 @@ class EnvGet(AbstractExpression):
 
         lookup_kind_expr = construct(self.lookup_kind, T.LookupKind)
 
-        ctx = get_context()
-
-        if self.categories:
-            check_source_language(
-                isinstance(self.categories, dict),
-                "Categories should be a dict"
+        # If no category is provided, consider they are all requested
+        if self.categories is None:
+            categories = RefCategories.Expr(get_context().ref_cats)
+        else:
+            categories = construct(
+                self.categories,
+                T.RefCategories,
+                'Invalid categories: {expected} expected but got {expr_type}'
             )
 
-            self.categories = {cat.lower(): val
-                               for cat, val in self.categories.items()}
-
-            check_source_language(
-                self.categories.get('others', None)
-                or all(self.categories.get(cat, None) for cat in ctx.ref_cats),
-                'Categories for env.get do not contain mappings for all'
-                ' categories'
-            )
-
-            check_source_language(
-                all(isinstance(val, bool) for val in self.categories.values()),
-                "Categories values should be bool"
-            )
-
-            self.categories = {
-                names.Name.from_lower(cat).camel_with_underscores
-                if cat != 'others' else 'others': val
-                for cat, val in self.categories.items()
-            }
-
-        return EnvGet.Expr(env_expr, sym_expr, lookup_kind_expr,
+        return EnvGet.Expr(env_expr, sym_expr, lookup_kind_expr, categories,
                            from_expr, self.only_first,
-                           abstract_expr=self, categories=self.categories)
+                           abstract_expr=self)
 
     def __repr__(self):
         return '<EnvGet({}, {})>'.format(self.env, self.symbol)

testsuite/tests/lexical_envs/categories/expected_concrete_syntax.lkt

@@ -0,0 +1,58 @@
+grammar None_grammar {
+    name <- Name(@Identifier)
+    main_rule <- Def(
+        @Def name @LBrace cat1 cat2 Example(@Example) @RBrace
+    )
+    cat1 <- Cat1(@Identifier("cat1") @LBrace decls @RBrace)
+    cat2 <- Cat2(@Identifier("cat2") @LBrace decls @RBrace)
+    var <- Var(name)
+    decls <- list+(var)
+
+}
+
+class FooNode {
+}
+
+class Cat1 : FooNode {
+    @parse_field decls : ASTList[Var]
+}
+
+class Cat2 : FooNode {
+    @parse_field decls : ASTList[Var]
+}
+
+class Def : FooNode {
+    @parse_field name : Name
+    @parse_field cat1 : Cat1
+    @parse_field cat2 : Cat2
+    @parse_field example : Example
+}
+
+class Example : FooNode {
+
+    @export fun lookup_all (name : SymbolType): Array[Entity[FooNode]] =
+    self.children_env.get(name)
+
+    @export fun lookup_1 (name : SymbolType): Array[Entity[FooNode]] =
+    self.children_env.get(
+        name, categories=RefCats(cat1=true, others=false)
+    )
+
+    @export fun lookup_2 (name : SymbolType): Array[Entity[FooNode]] =
+    self.children_env.get(
+        name, categories=RefCats(cat2=true, others=false)
+    )
+}
+
+"""
+List of Var.
+"""
+class ASTList[Var] : FooNodeBaseList {
+}
+
+class Name : FooNode {
+}
+
+class Var : FooNode {
+    @parse_field name : Name
+}

testsuite/tests/lexical_envs/categories/main.py

@@ -0,0 +1,38 @@
+from __future__ import absolute_import, division, print_function
+
+import sys
+
+import libfoolang
+
+
+print('main.py: Running...')
+
+
+ctx = libfoolang.AnalysisContext()
+u = ctx.get_from_buffer('input', """
+def a {
+    cat1 { b c   }
+    cat2 { b   d }
+    example
+}
+""")
+if u.diagnostics:
+    for d in u.diagnostics:
+        print(d)
+    sys.exit(1)
+
+
+def var_image(v):
+    return '<Var {}, line {}>'.format(v.text, v.sloc_range.start.line)
+
+
+example = u.root[-1]
+for lookup in ('p_lookup_all', 'p_lookup_1', 'p_lookup_2'):
+    print('With {}:'.format(lookup))
+    for name in ('a', 'b', 'c', 'd'):
+        lookup_prop = getattr(example, lookup)
+        print('  ({}) -> {}'.format(name,
+                                    [var_image(v) for v in lookup_prop(name)]))
+    print('')
+
+print('main.py: Done.')

testsuite/tests/lexical_envs/categories/test.out

@@ -0,0 +1,21 @@
+main.py: Running...
+With p_lookup_all:
+  (a) -> []
+  (b) -> ['<Var b, line 3>', '<Var b, line 4>']
+  (c) -> ['<Var c, line 3>']
+  (d) -> ['<Var d, line 4>']
+
+With p_lookup_1:
+  (a) -> []
+  (b) -> ['<Var b, line 3>']
+  (c) -> ['<Var c, line 3>']
+  (d) -> []
+
+With p_lookup_2:
+  (a) -> []
+  (b) -> ['<Var b, line 4>']
+  (c) -> []
+  (d) -> ['<Var d, line 4>']
+
+main.py: Done.
+Done

testsuite/tests/lexical_envs/categories/test.py

@@ -0,0 +1,92 @@
+"""
+Test that categories are properly considered during lexical env lookups.
+"""
+
+from __future__ import absolute_import, division, print_function
+
+from langkit.dsl import ASTNode, Field, T
+from langkit.envs import EnvSpec, add_env, add_to_env_kv, reference
+from langkit.expressions import RefCategories, Self, langkit_property
+from langkit.parsers import Grammar, List
+
+from lexer_example import Token
+from utils import build_and_run
+
+
+class FooNode(ASTNode):
+    pass
+
+
+class Def(FooNode):
+    name = Field(type=T.Name)
+    cat1 = Field(type=T.Cat1)
+    cat2 = Field(type=T.Cat2)
+    example = Field(type=T.Example)
+
+    env_spec = EnvSpec(add_env())
+
+
+class Cat1(FooNode):
+    decls = Field(type=T.Var.list)
+
+    env_spec = EnvSpec(add_env())
+
+
+class Cat2(FooNode):
+    decls = Field(type=T.Var.list)
+
+    env_spec = EnvSpec(add_env())
+
+
+class Example(FooNode):
+    token_node = True
+
+    env_spec = EnvSpec(
+        add_env(),
+        reference([Self.parent.cast(Def).cat1.cast(T.FooNode)],
+                  T.FooNode.children_env,
+                  category='cat1'),
+        reference([Self.parent.cast(Def).cat2.cast(T.FooNode)],
+                  T.FooNode.children_env,
+                  category='cat2')
+    )
+
+    @langkit_property(public=True)
+    def lookup_all(name=T.Symbol):
+        return Self.children_env.get(name)
+
+    @langkit_property(public=True)
+    def lookup_1(name=T.Symbol):
+        return Self.children_env.get(name, categories=RefCategories(cat1=True))
+
+    @langkit_property(public=True)
+    def lookup_2(name=T.Symbol):
+        return Self.children_env.get(name, categories=RefCategories(cat2=True))
+
+
+class Name(FooNode):
+    token_node = True
+
+
+class Var(FooNode):
+    name = Field(type=T.Name)
+
+    env_spec = EnvSpec(add_to_env_kv(Self.name.symbol, Self))
+
+
+G = Grammar('main_rule')
+G.add_rules(
+    main_rule=Def('def', G.name, '{',
+                  G.cat1,
+                  G.cat2,
+                  Example('example'),
+                  '}'),
+    cat1=Cat1(Token.Identifier('cat1'), '{', G.decls, '}'),
+    cat2=Cat2(Token.Identifier('cat2'), '{', G.decls, '}'),
+    decls=List(G.var, empty_valid=True),
+    var=Var(G.name),
+    name=Name(Token.Identifier),
+)
+
+build_and_run(G, 'main.py')
+print('Done')

testsuite/tests/lexical_envs/categories/test.yaml

@@ -0,0 +1 @@
+driver: python