scons-plus-plus/addons/astgen.py

import gzip
import json
import os.path
import pickle
import subprocess
from abc import ABC, abstractmethod
from typing import Callable, Any, Iterable, Self, Generator

from SCons.Script import *
from SCons.Node.FS import File
from spp import get_spp

spp = get_spp()

def post_environment(**kwargs) -> None:
    env: Environment = spp.globals['env']

    ast_json_builder = Builder(
        action=_gen_ast_json
    )
    env.Append(BUILDERS = {'AstJson': ast_json_builder})
    # env.SetDefault(ASTJSONCOM = '$ASTJSON -Xclang -ast-dump=json -fsyntax-only -Wno-unknown-warning-option -DSPP_AST_GEN $CXXFLAGS $SOURCES > $TARGET')

    env.AddMethod(_ast_jinja, 'AstJinja')

def _gen_ast_json(target: list[File], source: list[File], env: Environment):
    clang_exe = env.WhereIs('clang++')
    cmd = [clang_exe, '-Xclang', '-ast-dump=json', '-fsyntax-only', '-Wno-unknown-warning-option',
           '-DSPP_AST_GEN', f'-std={env["CXX_STANDARD"]}']
    for define in env['CPPDEFINES']:
        cmd.append(f'-D{define}')
    for path in env['CPPPATH']:
        cmd.append(f'-I{path}')
    cmd.append(source[0].abspath)
    # print(*cmd)
    try:
        proc = subprocess.Popen(cmd, text=True, stdout=subprocess.PIPE)
    except subprocess.CalledProcessError as e:
        env.Error(f'Clang exited with code {e.returncode}.')
        return
    parsed = json.load(proc.stdout)
    inner: list = parsed["inner"]

    # pos = 0
    # last_file = None

    #while pos < len(inner):
    #    last_file = inner[pos]["loc"].get("file", last_file)
    #    if last_file is None: #  or os.path.isabs(last_file):
    #        del inner[pos]
    #    else:
    #        pos += 1

    if target[0].suffix == '.bin':
        with gzip.open(target[0].abspath, 'wb') as f:
            pickle.dump(parsed, f)
    elif target[0].suffix == '.gz':
        with gzip.open(target[0].abspath, 'wt') as f:
            json.dump(parsed, f)
    else:
        with open(target[0].abspath, 'wt') as f:
            json.dump(parsed, f)

class ASTNode(ABC):
    @abstractmethod
    def _get_decls(self) -> Iterable[dict]: ...

    def inner(self) -> Iterable[dict]:
        return itertools.chain(*(decl['inner'] for decl in self._get_decls()))

    def inner_filtered(self, **kwargs) -> Iterable[dict]:
        def _applies(decl: dict) -> bool:
            for name, val in kwargs.items():
                if decl.get(name) != val:
                    return False
            return True
        return (decl for decl in self.inner() if _applies(decl))

class SimpleASTNode(ASTNode):
    def __init__(self, decl: dict) -> None:
        self._decl = decl

    def _get_decls(self) -> Iterable[dict]:
        return (self._decl,)

class Value(SimpleASTNode): ...

class Annotation(SimpleASTNode):
    @property
    def values(self) -> Iterable[Value]:
        return (Value(decl) for decl in self.inner())

class Param(SimpleASTNode):
    @property
    def name(self) -> str:
        return self._decl.get('name', '')

    @property
    def type(self) -> str:
        return self._decl['type']['qualType']

class Method(SimpleASTNode):
    def __init__(self, decl: dict, access: str) -> None:
        super().__init__(decl)
        self._access = access

    @property
    def access(self) -> str:
        return self._access

    @property
    def name(self) -> str:
        return self._decl['name']

    @property
    def mangled_name(self) -> str:
        return self._decl['mangledName']

    @property
    def type(self) -> str:
        return self._decl['type']['qualType']

    @property
    def return_type(self) -> str:
        return self.type.split('(', 1)[0].strip()

    @property
    def params(self) -> Iterable[Param]:
        return (Param(decl) for decl in self.inner_filtered(kind='ParmVarDecl'))

    @property
    def annotations(self) -> Iterable[Annotation]:
        return (Annotation(decl) for decl in self.inner_filtered(kind='AnnotateAttr'))

class Class(SimpleASTNode):
    @property
    def name(self) -> str:
        return self._decl['name']

    @property
    def tagUsed(self) -> str:
        return self._decl['tagUsed']

    @property
    def methods(self) -> Generator[Method]:
        access = 'private' if self.tagUsed == 'class' else 'public'
        for decl in self.inner():
            if decl['kind'] == 'AccessSpecDecl':
                access = decl['access']
            elif decl['kind'] == 'CXXMethodDecl' and not decl.get('isImplicit', False):
                yield Method(decl, access)

class Namespace(ASTNode, ABC):
    def get_namespace(self, ns_name: str) -> Self:
        return InnerNamespace(list(self.inner_filtered(kind='NamespaceDecl', name=ns_name)))

    @property
    def classes(self) -> Iterable[Class]:
        return (Class(decl) for decl in self.inner_filtered(kind='CXXRecordDecl', tagUsed='class', completeDefinition=True))

class InnerNamespace(Namespace):
    def __init__(self, decls: list[dict]) -> None:
        self._decls = decls

    def _get_decls(self) -> Iterable[dict]:
        return self._decls

class Ast(Namespace):
    def __init__(self, file: File) -> None:
        self._file = file
        self._data_dict: dict|None = None

    def _get_decls(self) -> tuple[dict]:
        if self._data_dict is None:
            if not self._file.exists():
                self._data_dict = {
                    'inner': []
                }
            elif self._file.suffix == '.bin':
                with gzip.open(self._file.abspath, 'rb') as f:
                    self._data_dict = pickle.load(f)
            elif self._file.suffix == '.gz':
                with gzip.open(self._file.abspath) as f:
                    self._data_dict = json.load(f)
            else:
                with open(self._file.abspath, 'r') as f:
                    self._data_dict = json.load(f)
        return (self._data_dict,)

def _ast_jinja(env: Environment, source: File, target: File, template: File, **kwargs):
    cache_dir = env['CACHE_DIR']
    rel_path = env.Dir('#').rel_path(source)
    json_file = env.File(os.path.join(cache_dir, 'ast_json', f'{rel_path}.bin'))
    ast_json = env.AstJson(target=json_file, source=source, **kwargs)

    ast_jinja = env.Jinja(
        target=target,
        source=template,
        JINJA_CONTEXT = {
            'ast': Ast(json_file)
        },
        **kwargs
    )
    env.Depends(ast_jinja, ast_json)
    # env.AlwaysBuild(ast_jinja)
    # env.Requires(ast_jinja, ast_json)
    # env.Requires(source, ast_jinja)
    env.Ignore(ast_json, ast_jinja)
    return ast_jinja