nvim temporary naga wgsl workaround

nvim/.config/nvim/bin/wgsl_flat.py

@@ -0,0 +1,546 @@
+#!/usr/bin/env python3
+"""Flatten naga_oil-style Bevy WGSL shaders into plain WGSL for wgsl-analyzer.
+
+Bevy preprocesses shaders with naga_oil (#import / #define_import_path /
+#ifdef / #{} shader-defs). wgsl-analyzer (>= 0.11.x) no longer understands
+that syntax, so these files are analyzed through flattened output generated by
+this script.
+
+Two modes:
+
+  wgsl_flat.py PROJECT [OUT_DIR]      twin mode (writes files)
+      Flattens every shader in PROJECT/assets/shaders into OUT_DIR
+      (default PROJECT/target/wgsl_flat). This is the on-disk debug path.
+
+  wgsl_flat.py --map FILE             map mode (writes JSON to stdout)
+      Flattens a single shader in memory and emits JSON with a line/column
+      map so the editor can project analyzer diagnostics back onto the real
+      file. See README comment in flatten_file_mapped for the layout.
+
+Resolution rules (naga_oil faithful):
+  #import a::b                -> whole module a::b
+  #import a::b::Item          -> item Item from module a::b
+  #import a::b::{X, Y}        -> items X, Y from module a::b
+  #import a::{b::X, c}        -> item X from a::b, whole module a::c
+  #import a::b as alias       -> module a::b, alias recorded for de-qualifying
+
+Modules are located by scanning crate sources for their #define_import_path.
+Qualified references through an alias name (alias::item) are rewritten to the
+bare item name because everything is inlined into one namespace.
+"""
+import glob
+import json
+import os
+import re
+import sys
+
+DEFINES = {"VERTEX_POSITIONS", "VERTEX_NORMALS", "VERTEX_UVS_A", "VERTEX_COLORS"}
+
+# naga_oil `#{NAME}` shader-def substitutions, desktop defaults from bevy_pbr.
+DEF_VALUES = {
+    "MATERIAL_BIND_GROUP": "3",
+    "MAX_DIRECTIONAL_LIGHTS": "1",
+    "MAX_CASCADES_PER_LIGHT": "4",
+    "MAX_RECT_LIGHTS": "8",
+    "MAX_POINT_LIGHTS": "8",
+    "MAX_SPOT_LIGHTS": "10",
+    "AVAILABLE_STORAGE_BUFFER_BINDINGS": "8",
+}
+
+IMPORT_RE = re.compile(r"^\s*#\s*import\s+(.*)$")
+DEFINE_IMPORT_RE = re.compile(r"^\s*#\s*define_import_path\s+.*$")
+IFDEF_RE = re.compile(r"^\s*#\s*ifdef\s+([\w]+)")
+IFNDEF_RE = re.compile(r"^\s*#\s*ifndef\s+([\w]+)")
+IF_RE = re.compile(r"^\s*#\s*if\s+")
+ELSE_RE = re.compile(r"^\s*#\s*else\b")
+ENDIF_RE = re.compile(r"^\s*#\s*endif\b")
+
+
+def subst_defs(line, pad=False):
+    """Replace #{NAME} shader-defs. With pad=True keep the original width so
+    column positions are preserved (used for diagnostic mapping)."""
+
+    def rep(m):
+        val = DEF_VALUES.get(m.group(1), "0")
+        if pad:
+            width = m.end() - m.start()
+            val = val[:width].ljust(width)
+        return val
+
+    return re.sub(r"#\{([A-Za-z_][\w]*)\}", rep, line)
+
+
+def extract_name(lines):
+    """Best-effort WGSL top-level item name from its header lines. `lines` is
+    a list of (lineno, text) tuples."""
+    text = " ".join(l[1].strip() for l in lines)
+    m = re.search(r"\bstruct\s+([A-Za-z_]\w*)", text)
+    if m:
+        return m.group(1)
+    m = re.search(r"\bfn\s+([A-Za-z_]\w*)", text)
+    if m:
+        return m.group(1)
+    m = re.search(r"\bconst\s+([A-Za-z_]\w*)", text)
+    if m:
+        return m.group(1)
+    m = re.search(r"\balias\s+([A-Za-z_]\w*)", text)
+    if m:
+        return m.group(1)
+    m = re.search(r"\blet\s+([A-Za-z_]\w*)", text)
+    if m:
+        return m.group(1)
+    m = re.search(r"\bvar(?:\s*<[^>]*>)?\s+([A-Za-z_]\w*)", text)
+    if m:
+        return m.group(1)
+    return None
+
+
+class Flattener:
+    def __init__(self, registry_dir, project_shaders=None):
+        self.modules = {}
+        self.sources = {}
+        self.visited = set()
+        self.qualifiers = set()
+        self.missing = set()
+        if project_shaders:
+            self.scan_dir(project_shaders)
+        for crate in ("bevy_pbr", "bevy_render"):
+            dirs = sorted(glob.glob(os.path.join(registry_dir, "*", crate) + "-*"))
+            if not dirs:
+                continue
+            self.scan_dir(os.path.join(dirs[-1], "src"))
+
+    def scan_dir(self, src_root):
+        for path in glob.glob(os.path.join(src_root, "**", "*.wgsl"),
+                              recursive=True):
+            try:
+                with open(path) as f:
+                    head = f.read(4096)
+            except OSError:
+                continue
+            key = None
+            for line in head.splitlines():
+                if DEFINE_IMPORT_RE.match(line):
+                    key = line.split(None, 2)[1].strip()
+                    break
+            if key is None:
+                continue
+            if key not in self.modules:
+                self.modules[key] = path
+                self.sources[path] = None
+
+    def module_path(self, key):
+        return self.modules.get(key)
+
+    def split_import(self, body):
+        """Return (module_path_or_None, [(module, item)], alias_or_None)."""
+        body = body.strip()
+        alias = None
+        if " as " in body:
+            body, alias = body.split(" as ", 1)
+            body = body.strip()
+            alias = alias.strip()
+        brace = body.find("::{")
+        if brace == -1:
+            return body, None, alias
+        prefix = body[:brace]
+        inner = body[brace + 3 :]
+        if inner.endswith("}"):
+            inner = inner[:-1]
+        entries = self.split_top_level(inner)
+        results = []
+        prefix_is_module = prefix in self.modules
+        for entry in entries:
+            entry = entry.strip()
+            if not entry:
+                continue
+            if "::{" in entry:
+                sub, _, subbody = entry.partition("::{")
+                module = prefix + "::" + sub.strip()
+                for name in self.split_top_level(subbody):
+                    if name.strip():
+                        results.append((module, name.strip()))
+            elif "::" in entry:
+                sub, _, item = entry.partition("::")
+                results.append((prefix + "::" + sub.strip(), item.strip() or None))
+            elif prefix_is_module:
+                results.append((prefix, entry))
+            else:
+                results.append((prefix + "::" + entry, None))
+        return None, results, alias
+
+    def split_top_level(self, text):
+        entries = []
+        depth = 0
+        cur = []
+        for ch in text:
+            if ch == "{":
+                depth += 1
+            elif ch == "}":
+                depth -= 1
+            if ch == "," and depth == 0:
+                entries.append("".join(cur))
+                cur = []
+            else:
+                cur.append(ch)
+        entries.append("".join(cur))
+        return entries
+
+    def resolve(self, module_or_item):
+        """Given `a::b` or `a::b::Item`, return (module, item_or_None)."""
+        key = module_or_item.strip()
+        if key in self.modules:
+            return key, None
+        head, _, tail = key.rpartition("::")
+        if tail and head in self.modules:
+            return head, tail
+        return None, None
+
+    def scan(self, text):
+        """Resolve directives; return (entries, imports) for one module.
+
+        entries is a list of (lineno, kind, text) where kind is "code" for an
+        active code line (already #{def}-substituted) and "blank" for a line
+        that must not produce tokens (a directive or a line inside an inactive
+        #ifdef branch). imports is a list of (dep, item_or_None, alias,
+        lineno)."""
+        # Scope stack of (parent_active, own_condition); active = both.
+        scopes = [(True, True)]
+        entries = []        # (lineno, kind, text)
+        imports = []        # (module, select, alias, lineno)
+        lines = text.splitlines()
+        i = 0
+        while i < len(lines):
+            line = lines[i]
+            lineno = i + 1
+            m_ifdef = IFDEF_RE.match(line)
+            m_ifndef = IFNDEF_RE.match(line)
+            m_if = IF_RE.match(line)
+            m_else = ELSE_RE.match(line)
+            m_endif = ENDIF_RE.match(line)
+            if m_ifdef or m_ifndef or m_if:
+                parent = scopes[-1][0] and scopes[-1][1]
+                cond = True
+                if m_ifdef:
+                    cond = m_ifdef.group(1) in DEFINES
+                elif m_ifndef:
+                    cond = m_ifndef.group(1) not in DEFINES
+                scopes.append((parent, cond))
+                entries.append((lineno, "blank", None))
+                i += 1
+                continue
+            if m_else:
+                parent, own = scopes[-1]
+                scopes[-1] = (parent, not own)
+                entries.append((lineno, "blank", None))
+                i += 1
+                continue
+            if m_endif:
+                if len(scopes) > 1:
+                    scopes.pop()
+                entries.append((lineno, "blank", None))
+                i += 1
+                continue
+            if not (scopes[-1][0] and scopes[-1][1]):
+                entries.append((lineno, "blank", None))
+                i += 1
+                continue
+            if DEFINE_IMPORT_RE.match(line):
+                entries.append((lineno, "blank", None))
+                i += 1
+                continue
+            m = IMPORT_RE.match(line)
+            if m:
+                body = m.group(1)
+                depth = body.count("{") - body.count("}")
+                start_lineno = lineno
+                while depth > 0 and i + 1 < len(lines):
+                    i += 1
+                    entries.append((i + 1, "blank", None))
+                    body += "\n" + lines[i]
+                    depth += lines[i].count("{") - lines[i].count("}")
+                mod, items, alias = self.split_import(body)
+                if mod is not None:
+                    dep, item = self.resolve(mod)
+                    if dep is None:
+                        self.missing.add(mod)
+                    else:
+                        imports.append((dep, item, alias, start_lineno))
+                else:
+                    for mod2, item2 in items:
+                        dep, _ = self.resolve(mod2)
+                        if dep is None:
+                            self.missing.add(mod2)
+                        else:
+                            imports.append((dep, item2, alias, start_lineno))
+                entries.append((start_lineno, "blank", None))
+                i += 1
+                continue
+            entries.append((lineno, "code", subst_defs(line)))
+            i += 1
+        return entries, imports
+
+    def split_items(self, kept):
+        """Group kept (lineno, text) lines into top-level WGSL items."""
+        items = []
+        current = []
+        depth = 0
+        prev_closed = True
+        for lineno, line in kept:
+            stripped = line.strip()
+            if not stripped:
+                if current:
+                    current.append((lineno, line))
+                continue
+            commentless = stripped.split("//", 1)[0].rstrip()
+            if depth == 0 and current and prev_closed:
+                items.append(current)
+                current = []
+            current.append((lineno, line))
+            depth += line.count("{") - line.count("}")
+            if not stripped.startswith("//"):
+                prev_closed = commentless.endswith((";", "}"))
+        if current:
+            items.append(current)
+        return items
+
+    def code_entries(self, entries):
+        return [(ln, tx) for ln, kind, tx in entries if kind == "code"]
+
+    def visit(self, module, select, out):
+        """Append (path, lineno, text) triples for `module` into `out`."""
+        if isinstance(select, str):
+            select = {select}
+        path = self.modules.get(module)
+        if path is None:
+            return
+        if module in self.in_progress:
+            return
+        if self.emitted.get(module) is None and module in self.emitted:
+            return
+        self.in_progress.add(module)
+        if self.sources[path] is None:
+            with open(path) as f:
+                self.sources[path] = f.read()
+        entries, directive_imports = self.scan(self.sources[path])
+        items = self.split_items(self.code_entries(entries))
+        for dep_mod, dep_select, alias, _ in directive_imports:
+            self.qualifiers.add(dep_mod.rsplit("::", 1)[-1])
+            if alias:
+                self.qualifiers.add(alias)
+            self.visit(dep_mod, dep_select, out)
+        prev = self.emitted.get(module, set())
+        if select is None:
+            if prev is not None:
+                for group in items:
+                    out.extend((path, ln, tx) for ln, tx in group)
+                self.emitted[module] = None
+        else:
+            if prev is None:
+                self.in_progress.discard(module)
+                return
+            wanted = {n for n in select if n not in prev}
+            for group in self.extract_with_deps(items, wanted):
+                name = extract_name(group)
+                if name not in prev:
+                    out.extend((path, ln, tx) for ln, tx in group)
+                    prev.add(name)
+            self.emitted[module] = prev
+        self.in_progress.discard(module)
+
+    def extract_with_deps(self, items, select):
+        """Items in select plus same-module top-level items they reference."""
+        by_name = {}
+        for group in items:
+            name = extract_name(group)
+            if name:
+                by_name[name] = group
+        want = set(select)
+        result = []
+        done = set()
+        while want:
+            name = want.pop()
+            if name in done or name not in by_name:
+                continue
+            done.add(name)
+            group = by_name[name]
+            result.append(group)
+            text = "\n".join(tx for _, tx in group)
+            for other in by_name:
+                if other not in done and re.search(r"\b" + re.escape(other) + r"\b", text):
+                    want.add(other)
+        return result
+
+    def _reset(self):
+        self.emitted = {}
+        self.in_progress = set()
+        self.qualifiers = set()
+        self.missing = set()
+
+    def flatten_file(self, path):
+        """Twin mode: return flattened plain-WGSL text (original behavior)."""
+        self._reset()
+        with open(path) as f:
+            text = f.read()
+        entries, imports = self.scan(text)
+        items = self.split_items(self.code_entries(entries))
+        out = []
+        for dep, select, alias, _ in imports:
+            self.qualifiers.add(dep.rsplit("::", 1)[-1])
+            if alias:
+                self.qualifiers.add(alias)
+            self.visit(dep, select, out)
+        for group in items:
+            out.extend((path, ln, tx) for ln, tx in group)
+        flat = "\n".join(tx for _, _, tx in out)
+        for q in sorted(self.qualifiers, key=len, reverse=True):
+            flat = re.sub(r"\b" + re.escape(q) + r"::", "", flat)
+        return flat
+
+    # ---- map mode ---------------------------------------------------------
+
+    @staticmethod
+    def _apply_sub(line, colmap, pat):
+        """Apply one regex removal pass, updating colmap (current -> original
+        column). Mirrors re.sub semantics."""
+        out = []
+        newmap = []
+        pos = 0
+        for m in re.finditer(pat, line):
+            s, e = m.start(), m.end()
+            out.append(line[pos:s])
+            newmap.extend(colmap[pos:s])
+            pos = e
+        out.append(line[pos:])
+        newmap.extend(colmap[pos:len(line) + 1])
+        return "".join(out), newmap
+
+    @classmethod
+    def _strip_with_map(cls, line, qualifiers):
+        """Strip qualifier:: prefixes, returning (new_line, colmap_or_None).
+        colmap[i] is the original column for flattened position i."""
+        colmap = list(range(len(line) + 1))
+        cur = line
+        changed = False
+        for q in qualifiers:
+            pat = r"\b" + re.escape(q) + r"::"
+            if not re.search(pat, cur):
+                continue
+            cur, colmap = cls._apply_sub(cur, colmap, pat)
+            changed = True
+        return cur, (colmap if changed else None)
+
+    def flatten_file_mapped(self, path):
+        """Map mode: flatten in memory and return a JSON-serialisable dict.
+
+        Layout of the flattened document (0-based line indices):
+          [0, main_lines)       real file, 1:1 line correspondence
+          main_lines            single blank separator line (unmapped)
+          (main_lines, ...)     appended imported content; `extra[i]` gives
+                                provenance for flattened line main_lines+1+i
+
+        Returns {file, main_lines, code, colmaps, extra}:
+          colmaps: {flattened_line: [orig_col, ...]} only for lines where
+                   qualifier stripping changed column positions.
+          extra:   [null | [src_path, src_lineno, import_lineno], ...]
+        """
+        self._reset()
+        with open(path) as f:
+            text = f.read()
+        entries, imports = self.scan(text)
+        main_lines = len(text.splitlines())
+
+        # Main region: preserve every line. Directives and inactive branches
+        # become empty; active code is taken from the raw source with padded
+        # #{def} substitution so columns line up with the real file. (scan()
+        # substitutes unpadded for the twin path, so use the raw lines here.)
+        raw = text.splitlines()
+        kinds = {}
+        for lineno, kind, _ in entries:
+            kinds.setdefault(lineno, kind)
+        region = []
+        for ln in range(1, main_lines + 1):
+            if kinds.get(ln) == "code":
+                region.append(subst_defs(raw[ln - 1], pad=True))
+            else:
+                region.append("")
+
+        # Imported content, appended after the main region.
+        appended = []   # (src_path, src_lineno, import_lineno, text)
+        for dep, select, alias, import_lineno in imports:
+            self.qualifiers.add(dep.rsplit("::", 1)[-1])
+            if alias:
+                self.qualifiers.add(alias)
+            buf = []
+            self.visit(dep, select, buf)
+            for src_path, src_ln, tx in buf:
+                appended.append((src_path, src_ln, import_lineno, tx))
+
+        quals = sorted(self.qualifiers, key=len, reverse=True)
+        flat_lines = []
+        colmaps = {}
+        extra = []
+
+        for idx, line in enumerate(region):
+            new, cm = self._strip_with_map(line, quals)
+            flat_lines.append(new)
+            if cm is not None:
+                colmaps[idx] = cm
+
+        flat_lines.append("")  # separator, unmapped
+
+        for src_path, src_ln, import_ln, line in appended:
+            new, cm = self._strip_with_map(line, quals)
+            flat_lines.append(new)
+            li = len(flat_lines) - 1
+            if cm is not None:
+                colmaps[li] = cm
+            extra.append([src_path, src_ln, import_ln])
+
+        return {
+            "file": path,
+            "main_lines": main_lines,
+            "code": "\n".join(flat_lines),
+            "colmaps": {str(k): v for k, v in colmaps.items()},
+            "extra": extra,
+        }
+
+
+def run_twins(project, out_dir):
+    home = os.path.expanduser("~")
+    registry = os.path.join(home, ".cargo", "registry", "src")
+    flattener = Flattener(registry, os.path.join(project, "assets", "shaders"))
+    shaders_dir = os.path.join(project, "assets", "shaders")
+    os.makedirs(out_dir, exist_ok=True)
+    for shader in sorted(glob.glob(os.path.join(shaders_dir, "*.wgsl"))):
+        name = os.path.basename(shader)
+        flat = flattener.flatten_file(shader)
+        with open(os.path.join(out_dir, name), "w") as f:
+            f.write(flat)
+        tail = f", missing: {sorted(flattener.missing)}" if flattener.missing else ""
+        print(f"{name}: {len(flat.splitlines())} lines{tail}")
+
+
+def run_map(file_path):
+    home = os.path.expanduser("~")
+    registry = os.path.join(home, ".cargo", "registry", "src")
+    shaders_dir = os.path.dirname(os.path.abspath(file_path))
+    flattener = Flattener(registry, shaders_dir)
+    result = flattener.flatten_file_mapped(os.path.abspath(file_path))
+    if flattener.missing:
+        result["missing"] = sorted(flattener.missing)
+    json.dump(result, sys.stdout)
+    sys.stdout.write("\n")
+
+
+def main():
+    if len(sys.argv) >= 3 and sys.argv[1] == "--map":
+        run_map(sys.argv[2])
+        return
+    project = sys.argv[1] if len(sys.argv) > 1 else "."
+    out_dir = sys.argv[2] if len(sys.argv) > 2 else os.path.join(project, "target", "wgsl_flat")
+    run_twins(project, out_dir)
+
+
+if __name__ == "__main__":
+    main()

nvim/.config/nvim/init.lua

@@ -24,3 +24,4 @@ require("bearded-arc.watcher").setup()
 require "options"
 require "autocmds"
 require "mappings"
+require "configs.wgsl_lsp"

nvim/.config/nvim/lua/configs/lspconfig.lua

@@ -61,6 +61,23 @@ local servers = {
   },
   wgsl_analyzer = {
     filetypes = { "wgsl", "wesl" },
+    -- Real naga_oil shaders (assets/shaders/*.wgsl) use `#import`/`#{}`
+    -- syntax that wgsl-analyzer can't parse. Skip them entirely -- never
+    -- calling on_dir() means the client does not attach. Diagnostics for
+    -- those files run on in-memory shadow buffers instead (see
+    -- configs.wgsl_lsp) and are projected back onto the real files.
+    -- Shadow buffers are buftype=nofile, so vim.lsp.enable skips them and
+    -- wgsl_lsp starts the client for them explicitly -- no special-casing
+    -- needed here. Everything else attaches normally; outside a git root it
+    -- falls back to single-file mode.
+    root_dir = function(bufnr, on_dir)
+      local name = vim.api.nvim_buf_get_name(bufnr)
+      if name:match("assets/shaders/[^/]+%.wgsl$") then
+        return
+      end
+      local util = require("lspconfig.util")
+      on_dir(util.root_pattern(".git")(name))
+    end,
     on_attach = on_attach,
     capabilities = capabilities,
   },
@@ -82,16 +99,9 @@ local servers = {
   },
 }
 
+-- vim.lsp.enable() reads `filetypes` from each config and sets up its own
+-- FileType autocmds, so no extra autocmds are needed here.
 for server, opts in pairs(servers) do
   vim.lsp.config(server, opts)
-  if opts.filetypes then
+  vim.lsp.enable(server)
-    vim.api.nvim_create_autocmd("FileType", {
-      pattern = opts.filetypes,
-      callback = function()
-        vim.lsp.enable(server)
-      end,
-    })
-  else
-    vim.lsp.enable(server)
-  end
 end

nvim/.config/nvim/lua/configs/wgsl_lsp.lua

@@ -0,0 +1,298 @@
+-- naga_oil (Bevy) shaders are not plain WGSL: wgsl-analyzer cannot parse the
+-- `#import` / `#ifdef` / `#{}` shader-def syntax, so no LSP ever attaches to
+-- the real shader files (see the root_dir() gate in lua/configs/lspconfig.lua).
+--
+-- Instead each shader gets an unlisted "shadow" buffer holding a flattened
+-- plain-WGSL rendering produced by bin/wgsl_flat.py --map. wgsl-analyzer is
+-- started on the shadow buffer explicitly (vim.lsp.enable skips buftype=
+-- nofile buffers), and its diagnostics are projected back onto the real file:
+--
+--   * real-file lines map 1:1 into the flattened document (directives are
+--     blanked, #{} values padded to keep widths), so diagnostics land on the
+--     exact line and column;
+--   * diagnostics inside inlined imports are attributed to the #import line
+--     that pulled them in, prefixed with the origin file and line.
+--
+-- Nothing is written to disk.
+--
+--   :WgslLsp      open the shadow buffer of the current shader in a vsplit
+--                 (full LSP features on the flattened view)
+--   :WgslFlatten  write on-disk twins into <project>/target/wgsl_flat/
+--                 (debug aid only; the LSP no longer needs them)
+
+local M = {}
+
+local SCRIPT = vim.fn.stdpath("config") .. "/bin/wgsl_flat.py"
+local NS = vim.api.nvim_create_namespace("wgsl_flat")
+local MARK = "/.wgsl_lsp/"
+
+-- abs shader path -> { file, flat_buf, seq, main_lines, colmaps, extra }
+M.sessions = {}
+
+local grp = vim.api.nvim_create_augroup("wgsl_flat", { clear = true })
+
+local function project_root(file)
+  local cargo = vim.fn.findfile("Cargo.toml", file .. ";")
+  if cargo == "" then
+    return nil
+  end
+  return vim.fn.fnamemodify(cargo, ":h")
+end
+
+local function is_real_shader(name)
+  return name:match("assets/shaders/[^/]+%.wgsl$") ~= nil
+end
+
+local function new_session(path)
+  local project = project_root(path)
+  if not project then
+    return nil
+  end
+  local buf = vim.api.nvim_create_buf(false, true)
+  vim.api.nvim_buf_set_name(buf, project .. MARK .. vim.fn.fnamemodify(path, ":t"))
+  -- filetype wgsl for highlighting if the buffer is ever displayed; the LSP
+  -- client is started explicitly by attach_client (vim.lsp.enable skips
+  -- buftype=nofile buffers).
+  vim.bo[buf].filetype = "wgsl"
+  local sess = { file = path, flat_buf = buf, seq = 0 }
+  M.sessions[path] = sess
+  return sess
+end
+
+-- Attach wgsl-analyzer to the shadow buffer. vim.lsp.enable never attaches
+-- it on its own: the shadow buffer is buftype=nofile, which the enable
+-- callback skips. So start the client explicitly from the merged
+-- vim.lsp.config. Retries because nvim-lspconfig lazy-loads on User
+-- FilePost, which may happen after the first refresh.
+local function attach_client(sess, tries)
+  tries = tries or 0
+  if not vim.api.nvim_buf_is_valid(sess.flat_buf) then
+    return
+  end
+  if #vim.lsp.get_clients({ bufnr = sess.flat_buf }) > 0 then
+    return
+  end
+  local c = vim.lsp.config["wgsl_analyzer"]
+  if not c or not c.cmd then
+    if tries < 40 then
+      vim.defer_fn(function()
+        attach_client(sess, tries + 1)
+      end, 250)
+    end
+    return
+  end
+  local ok, err = pcall(vim.lsp.start, {
+    name = "wgsl_analyzer",
+    cmd = c.cmd,
+    root_dir = project_root(sess.file) or vim.fn.getcwd(),
+    capabilities = c.capabilities,
+    on_attach = c.on_attach,
+    settings = c.settings,
+  }, { bufnr = sess.flat_buf })
+  if not ok then
+    vim.notify("wgsl: could not start wgsl-analyzer: " .. tostring(err), vim.log.levels.ERROR)
+  end
+end
+
+vim.api.nvim_create_autocmd("User", {
+  group = grp,
+  pattern = "FilePost",
+  once = true,
+  callback = function()
+    for _, sess in pairs(M.sessions) do
+      attach_client(sess)
+    end
+  end,
+})
+
+--- Re-flatten `path` and push the result into its shadow buffer.
+function M.refresh(path, cb)
+  local sess = M.sessions[path] or new_session(path)
+  if not sess then
+    return
+  end
+  sess.seq = sess.seq + 1
+  local seq = sess.seq
+  vim.system({ "python3", SCRIPT, "--map", path }, function(res)
+    vim.schedule(function()
+      if not vim.api.nvim_buf_is_valid(sess.flat_buf) then
+        return
+      end
+      if res.code ~= 0 then
+        vim.notify(
+          "wgsl_flat failed:\n" .. ((res.stderr or "") .. (res.stdout or "")),
+          vim.log.levels.ERROR
+        )
+        return
+      end
+      if sess.seq ~= seq then
+        return -- superseded by a newer refresh
+      end
+      local ok, data = pcall(vim.json.decode, res.stdout)
+      if not ok or type(data) ~= "table" then
+        vim.notify("wgsl_flat: could not parse JSON output", vim.log.levels.ERROR)
+        return
+      end
+      sess.main_lines = data.main_lines or 0
+      sess.colmaps = data.colmaps or {}
+      sess.extra = data.extra or {}
+      local lines = vim.split(data.code or "", "\n", { plain = true })
+      vim.api.nvim_buf_set_lines(sess.flat_buf, 0, -1, false, lines)
+      attach_client(sess)
+      if cb then
+        pcall(cb)
+      end
+    end)
+  end)
+end
+
+local function session_by_flat(buf)
+  for _, sess in pairs(M.sessions) do
+    if sess.flat_buf == buf then
+      return sess
+    end
+  end
+end
+
+-- Flattened column -> real-file column for one line (qualifier strips).
+local function adj_col(sess, line, col)
+  local cm = sess.colmaps[tostring(line)]
+  if not cm or col >= #cm then
+    return col
+  end
+  return cm[col + 1] or col
+end
+
+--- Project shadow-buffer diagnostics back onto the real file.
+local function remap(sess, diags)
+  local out = {}
+  for _, d in ipairs(diags) do
+    if d.lnum < sess.main_lines then
+      -- main region: 1:1 line correspondence
+      local nd = vim.deepcopy(d)
+      nd.col = adj_col(sess, d.lnum, d.col or 0)
+      if d.end_lnum == d.lnum and d.end_col then
+        nd.end_col = adj_col(sess, d.lnum, d.end_col)
+      else
+        nd.end_lnum = d.lnum
+        nd.end_col = nil
+      end
+      table.insert(out, nd)
+    elseif d.lnum > sess.main_lines then
+      -- appended imports: attribute to the #import line that pulled them in
+      local e = sess.extra[d.lnum - sess.main_lines]
+      if e and e[3] then
+        table.insert(out, {
+          lnum = e[3] - 1,
+          col = 0,
+          severity = d.severity,
+          message = string.format(
+            "[%s:%d] %s",
+            vim.fn.fnamemodify(e[1] or "", ":t"),
+            e[2] or 0,
+            d.message
+          ),
+          source = d.source,
+        })
+      end
+    end
+  end
+  return out
+end
+
+local function publish(sess)
+  local mapped = remap(sess, vim.diagnostic.get(sess.flat_buf))
+  for _, buf in ipairs(vim.api.nvim_list_bufs()) do
+    if
+      vim.api.nvim_buf_is_loaded(buf)
+      and vim.api.nvim_buf_get_name(buf) == sess.file
+    then
+      vim.diagnostic.set(NS, buf, mapped)
+    end
+  end
+end
+
+vim.api.nvim_create_autocmd("DiagnosticChanged", {
+  group = grp,
+  callback = function(args)
+    local sess = session_by_flat(args.buf)
+    if sess then
+      publish(sess)
+    end
+  end,
+})
+
+-- Initial flatten when a real shader is opened, re-flatten on save. The
+-- shadow buffer content change feeds didChange to the attached client.
+vim.api.nvim_create_autocmd({ "BufReadPost", "BufWritePost" }, {
+  group = grp,
+  pattern = "*/assets/shaders/*.wgsl",
+  callback = function(args)
+    local name = vim.api.nvim_buf_get_name(args.buf)
+    if name ~= "" then
+      M.refresh(name)
+    end
+  end,
+})
+
+-- Drop the session when the last buffer of the real file goes away.
+vim.api.nvim_create_autocmd("BufWipeout", {
+  group = grp,
+  pattern = "*/assets/shaders/*.wgsl",
+  callback = function(args)
+    local name = vim.api.nvim_buf_get_name(args.buf)
+    for _, b in ipairs(vim.api.nvim_list_bufs()) do
+      if
+        b ~= args.buf
+        and vim.api.nvim_buf_is_loaded(b)
+        and vim.api.nvim_buf_get_name(b) == name
+      then
+        return
+      end
+    end
+    local sess = M.sessions[name]
+    if sess then
+      M.sessions[name] = nil
+      if vim.api.nvim_buf_is_valid(sess.flat_buf) then
+        vim.api.nvim_buf_delete(sess.flat_buf, { force = true })
+      end
+    end
+  end,
+})
+
+vim.api.nvim_create_user_command("WgslLsp", function()
+  local name = vim.api.nvim_buf_get_name(0)
+  if not is_real_shader(name) then
+    vim.notify("wgsl: current buffer is not a naga_oil shader", vim.log.levels.WARN)
+    return
+  end
+  M.refresh(name, function()
+    local sess = M.sessions[name]
+    if sess and vim.api.nvim_buf_is_valid(sess.flat_buf) then
+      vim.cmd("vsplit")
+      vim.api.nvim_set_current_buf(sess.flat_buf)
+    end
+  end)
+end, { desc = "Open the flattened shadow view of the current shader (wgsl-analyzer)" })
+
+vim.api.nvim_create_user_command("WgslFlatten", function()
+  local project = project_root(vim.api.nvim_buf_get_name(0))
+  if not project then
+    vim.notify("wgsl: not inside a Cargo project", vim.log.levels.WARN)
+    return
+  end
+  vim.system({ "python3", SCRIPT, project }, function(res)
+    vim.schedule(function()
+      if res.code ~= 0 then
+        vim.notify(
+          "wgsl_flat failed:\n" .. ((res.stderr or "") .. (res.stdout or "")),
+          vim.log.levels.ERROR
+        )
+      else
+        vim.notify("wgsl: twins written to " .. project .. "/target/wgsl_flat/")
+      end
+    end)
+  end)
+end, { desc = "Write flattened WGSL twins to target/wgsl_flat (debug aid)" })
+
+return M

nvim/.config/nvim/lua/plugins/coding.lua

@@ -33,10 +33,12 @@ return {
   },
 
   -- Treesitter parser manager
+  -- Loaded on startup: `cmd`-only laziness meant it never loaded in normal
+  -- sessions, breaking treesitter highlighting everywhere.
   {
     "nvim-treesitter/nvim-treesitter",
     build = ":TSUpdate",
-    cmd = { "TSInstall", "TSUpdate" },
+    event = "VeryLazy",
     config = function()
       require("nvim-treesitter").setup()
     end,