diff --git a/.envrc b/.envrc
new file mode 100644
index 0000000..720e019
--- /dev/null
+++ b/.envrc
@@ -0,0 +1,5 @@
+if ! has nix_direnv_version || ! nix_direnv_version 3.0.6; then
+  source_url "https://raw.githubusercontent.com/nix-community/nix-direnv/3.0.6/direnvrc" "sha256-RYcUJaRMf8oF5LznDrlCXbkOQrywm0HDv1VjYGaJGdM="
+fi
+
+use flake
diff --git a/.gitignore b/.gitignore
index e3fbd98..8e40704 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,2 +1,10 @@
+### direnv ###
+.direnv
+
+### nix ###
+result*
+repl-result-*
+
+### tree-sitter ###
 build
 node_modules
diff --git a/.yamlfmt b/.yamlfmt
new file mode 100644
index 0000000..e9d24f9
--- /dev/null
+++ b/.yamlfmt
@@ -0,0 +1,3 @@
+formatter:
+  type: basic
+  retain_line_breaks: true
diff --git a/default.nix b/default.nix
index 17d94ed..35191de 100644
--- a/default.nix
+++ b/default.nix
@@ -1,19 +1,23 @@
-{ pkgs ? import <nixpkgs> { }
-, lib ? pkgs.lib
-, src ? lib.cleanSource ./.
-}:
+{system ? builtins.currentSystem}: let
+  lock = builtins.fromJSON (builtins.readFile ./flake.lock);
 
-pkgs.tree-sitter-grammars.tree-sitter-nix.overrideAttrs (old: {
-  name = "tree-sitter-nix-dev";
-  version = "dev";
-  inherit src;
+  root = lock.nodes.${lock.root};
+  inherit
+    (lock.nodes.${root.inputs.flake-compat}.locked)
+    owner
+    repo
+    rev
+    narHash
+    ;
 
-  doCheck = true;
-  checkInputs = [
-    pkgs.tree-sitter
-    pkgs.nodejs
-  ];
-  checkPhase = ''
-    HOME=$(mktemp -d) tree-sitter test
-  '';
-})
+  flake-compat = fetchTarball {
+    url = "https://github.com/${owner}/${repo}/archive/${rev}.tar.gz";
+    sha256 = narHash;
+  };
+
+  flake = import flake-compat {
+    inherit system;
+    src = ./.;
+  };
+in
+  flake.defaultNix
diff --git a/devshell.nix b/devshell.nix
new file mode 100644
index 0000000..fb5193a
--- /dev/null
+++ b/devshell.nix
@@ -0,0 +1,16 @@
+{pkgs, ...}:
+pkgs.mkShell {
+  packages = with pkgs; [
+    nodejs
+    python3
+
+    tree-sitter
+    editorconfig-checker
+
+    rustc
+    cargo
+
+    go
+    gcc
+  ];
+}
diff --git a/flake.lock b/flake.lock
index 6dcc7e6..e8f0737 100644
--- a/flake.lock
+++ b/flake.lock
@@ -1,5 +1,20 @@
 {
   "nodes": {
+    "flake-compat": {
+      "locked": {
+        "lastModified": 1717312683,
+        "narHash": "sha256-FrlieJH50AuvagamEvWMIE6D2OAnERuDboFDYAED/dE=",
+        "owner": "nix-community",
+        "repo": "flake-compat",
+        "rev": "38fd3954cf65ce6faf3d0d45cd26059e059f07ea",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nix-community",
+        "repo": "flake-compat",
+        "type": "github"
+      }
+    },
     "flake-utils": {
       "inputs": {
         "systems": "systems"
@@ -56,9 +71,11 @@
     },
     "root": {
       "inputs": {
+        "flake-compat": "flake-compat",
         "flake-utils": "flake-utils",
         "nix-github-actions": "nix-github-actions",
-        "nixpkgs": "nixpkgs"
+        "nixpkgs": "nixpkgs",
+        "treefmt-nix": "treefmt-nix"
       }
     },
     "systems": {
@@ -75,6 +92,26 @@
         "repo": "default",
         "type": "github"
       }
+    },
+    "treefmt-nix": {
+      "inputs": {
+        "nixpkgs": [
+          "nixpkgs"
+        ]
+      },
+      "locked": {
+        "lastModified": 1727984844,
+        "narHash": "sha256-xpRqITAoD8rHlXQafYZOLvUXCF6cnZkPfoq67ThN0Hc=",
+        "owner": "numtide",
+        "repo": "treefmt-nix",
+        "rev": "4446c7a6fc0775df028c5a3f6727945ba8400e64",
+        "type": "github"
+      },
+      "original": {
+        "owner": "numtide",
+        "repo": "treefmt-nix",
+        "type": "github"
+      }
     }
   },
   "root": "root",
diff --git a/flake.nix b/flake.nix
index f1ebd25..15f1262 100644
--- a/flake.nix
+++ b/flake.nix
@@ -8,114 +8,119 @@
 
     nix-github-actions.url = "github:nix-community/nix-github-actions";
     nix-github-actions.inputs.nixpkgs.follows = "nixpkgs";
+
+    flake-compat.url = "github:nix-community/flake-compat";
+
+    treefmt-nix = {
+      url = "github:numtide/treefmt-nix";
+      inputs.nixpkgs.follows = "nixpkgs";
+    };
   };
 
-  outputs = { self, nixpkgs, flake-utils, nix-github-actions }: (
-    (flake-utils.lib.eachDefaultSystem (system:
-      let
-        pkgs = nixpkgs.legacyPackages.${system};
-        inherit (pkgs) lib;
+  outputs = {
+    self,
+    flake-utils,
+    nix-github-actions,
+    nixpkgs,
+    treefmt-nix,
+    ...
+  }: (
+    (flake-utils.lib.eachDefaultSystem (system: let
+      pkgs = nixpkgs.legacyPackages.${system};
+      inherit (pkgs) lib;
+      treefmtEval = treefmt-nix.lib.evalModule pkgs ./treefmt.nix;
+    in {
+      checks = let
+        # shellPackages = (pkgs.callPackage ./shell.nix { }).packages;
+        # If the generated code differs from the checked in we need
+        # to check in the newly generated sources.
+        mkCheck = name: check:
+          pkgs.runCommand name
+          {
+            inherit (self.devShells.${system}.default) nativeBuildInputs;
+          } ''
+            cp -rv ${self} src
+            chmod +w -R src
+            cd src
+
+            ${check}
 
+            touch $out
+          '';
       in
-      {
-        checks =
-          let
-            # shellPackages = (pkgs.callPackage ./shell.nix { }).packages;
-
-            # If the generated code differs from the checked in we need
-            # to check in the newly generated sources.
-            mkCheck = name: check: pkgs.runCommand name
-              {
-                inherit (self.devShells.${system}.default) nativeBuildInputs;
-              } ''
-              cp -rv ${self} src
-              chmod +w -R src
-              cd src
-
-              ${check}
-
-              touch $out
-            '';
+        {
+          build = self.packages.${system}.tree-sitter-nix;
+
+          editorconfig = mkCheck "editorconfig" "editorconfig-checker";
+
+          # If the generated code differs from the checked in we need
+          # to check in the newly generated sources.
+          generated-diff = mkCheck "generated-diff" ''
+            HOME=. npm run generate
+            diff -r src/ ${self}/src
+          '';
+
+          treefmt = treefmtEval.config.build.check self;
 
+          rust-bindings = let
+            cargo' = lib.importTOML ./Cargo.toml;
           in
-          {
-            build = self.packages.${system}.tree-sitter-nix;
-
-            editorconfig = mkCheck "editorconfig" "editorconfig-checker";
-
-            # If the generated code differs from the checked in we need
-            # to check in the newly generated sources.
-            generated-diff = mkCheck "generated-diff" ''
-              HOME=. npm run generate
-              diff -r src/ ${self}/src
-            '';
-
-            treefmt = mkCheck "treefmt" "treefmt --no-cache --fail-on-change";
-
-            rust-bindings =
-              let
-                cargo' = lib.importTOML ./Cargo.toml;
-              in
-              pkgs.rustPlatform.buildRustPackage {
-                pname = cargo'.package.name;
-                inherit (cargo'.package) version;
-                src = self;
-                cargoLock = {
-                  lockFile = ./Cargo.lock;
-                };
+            pkgs.rustPlatform.buildRustPackage {
+              pname = cargo'.package.name;
+              inherit (cargo'.package) version;
+              src = self;
+              cargoLock = {
+                lockFile = ./Cargo.lock;
               };
-
-          } // lib.optionalAttrs (!pkgs.stdenv.isDarwin) {
-            # Requires xcode
-            node-bindings =
-              let
-                package' = lib.importJSON ./package.json;
-              in
-              pkgs.stdenv.mkDerivation {
-                pname = package'.name;
-                inherit (package') version;
-                src = self;
-                nativeBuildInputs = with pkgs; [
-                  importNpmLock.hooks.npmConfigHook
-                  nodejs
-                  nodejs.passthru.python # for node-gyp
-                  npmHooks.npmBuildHook
-                  npmHooks.npmInstallHook
-                  tree-sitter
-                ];
-                npmDeps = pkgs.importNpmLock {
-                  npmRoot = ./.;
-                };
-                buildPhase = ''
-                  runHook preBuild
-                  ${pkgs.nodePackages.node-gyp}/bin/node-gyp configure
-                  npm run build
-                  runHook postBuild
-                '';
-                installPhase = "touch $out";
+            };
+        }
+        // lib.optionalAttrs (!pkgs.stdenv.isDarwin) {
+          # Requires xcode
+          node-bindings = let
+            package' = lib.importJSON ./package.json;
+          in
+            pkgs.stdenv.mkDerivation {
+              pname = package'.name;
+              inherit (package') version;
+              src = self;
+              nativeBuildInputs = with pkgs; [
+                importNpmLock.hooks.npmConfigHook
+                nodejs
+                nodejs.passthru.python # for node-gyp
+                npmHooks.npmBuildHook
+                npmHooks.npmInstallHook
+                tree-sitter
+              ];
+              npmDeps = pkgs.importNpmLock {
+                npmRoot = ./.;
               };
-          };
-
-        packages.tree-sitter-nix = pkgs.callPackage ./default.nix { src = self; };
-        packages.default = self.packages.${system}.tree-sitter-nix;
-        devShells.default = pkgs.callPackage ./shell.nix { };
+              buildPhase = ''
+                runHook preBuild
+                ${pkgs.nodePackages.node-gyp}/bin/node-gyp configure
+                npm run build
+                runHook postBuild
+              '';
+              installPhase = "touch $out";
+            };
+        };
 
-        formatter = pkgs.writeShellScriptBin "tree-sitter-nix-fmt" ''
-          exec ${pkgs.treefmt}/bin/treefmt --config-file ${./treefmt.toml} "$@"
-        '';
+      packages.tree-sitter-nix = pkgs.callPackage ./package.nix {src = self;};
+      packages.default = self.packages.${system}.tree-sitter-nix;
 
-      })) // {
+      devShells.default = pkgs.callPackage ./devshell.nix {};
 
+      formatter = treefmtEval.config.build.wrapper;
+    }))
+    // {
       githubActions = nix-github-actions.lib.mkGithubMatrix {
         # Inherit GHA actions matrix from a subset of platforms supported by hosted runners
         checks = {
           inherit (self.checks) x86_64-linux;
 
           # Don't run linters on darwin as it's just scheduling overhead
-          x86_64-darwin = builtins.removeAttrs self.checks.x86_64-darwin [ "editorconfig" "generated-diff" "treefmt" ];
+          x86_64-darwin = builtins.removeAttrs self.checks.x86_64-darwin ["editorconfig" "generated-diff" "treefmt"];
         };
       };
-
     }
   );
 }
diff --git a/grammar.js b/grammar.js
index d241a4d..2eadb9c 100644
--- a/grammar.js
+++ b/grammar.js
@@ -206,7 +206,7 @@ module.exports = grammar({
           prec(
             precedence,
             seq(field("operator", operator), field("argument", $._expr_op)),
-          ),
+          )
         ),
       ),
 
@@ -234,7 +234,7 @@ module.exports = grammar({
               field("operator", operator),
               field("right", $._expr_op),
             ),
-          ),
+          )
         ),
         // right assoc.
         ...[
@@ -249,7 +249,7 @@ module.exports = grammar({
               field("operator", operator),
               field("right", $._expr_op),
             ),
-          ),
+          )
         ),
       ),
 
diff --git a/package.nix b/package.nix
new file mode 100644
index 0000000..190fb5c
--- /dev/null
+++ b/package.nix
@@ -0,0 +1,19 @@
+{
+  pkgs ? import <nixpkgs> {},
+  lib ? pkgs.lib,
+  src ? lib.cleanSource ./.,
+}:
+pkgs.tree-sitter-grammars.tree-sitter-nix.overrideAttrs (_old: {
+  name = "tree-sitter-nix-dev";
+  version = "dev";
+  inherit src;
+
+  doCheck = true;
+  checkInputs = [
+    pkgs.tree-sitter
+    pkgs.nodejs
+  ];
+  checkPhase = ''
+    HOME=$(mktemp -d) tree-sitter test
+  '';
+})
diff --git a/renovate.json b/renovate.json
index f4f939c..5afba06 100644
--- a/renovate.json
+++ b/renovate.json
@@ -1,9 +1,13 @@
 {
   "$schema": "https://docs.renovatebot.com/renovate-schema.json",
-  "extends": ["config:recommended"],
+  "extends": [
+    "config:recommended"
+  ],
   "lockFileMaintenance": {
     "enabled": true,
-    "extends": ["schedule:weekly"]
+    "extends": [
+      "schedule:weekly"
+    ]
   },
   "nix": {
     "enabled": true
diff --git a/shell.nix b/shell.nix
index bce0d8f..bf35524 100644
--- a/shell.nix
+++ b/shell.nix
@@ -1,21 +1,23 @@
-{ pkgs ? import <nixpkgs> { } }:
+{system ? builtins.currentSystem}: let
+  lock = builtins.fromJSON (builtins.readFile ./flake.lock);
 
-pkgs.mkShell {
-  packages = [
-    pkgs.nodejs
-    pkgs.python3
+  root = lock.nodes.${lock.root};
+  inherit
+    (lock.nodes.${root.inputs.flake-compat}.locked)
+    owner
+    repo
+    rev
+    narHash
+    ;
 
-    pkgs.tree-sitter
-    pkgs.editorconfig-checker
+  flake-compat = fetchTarball {
+    url = "https://github.com/${owner}/${repo}/archive/${rev}.tar.gz";
+    sha256 = narHash;
+  };
 
-    pkgs.rustc
-    pkgs.cargo
-
-    # Formatters
-    pkgs.treefmt
-    pkgs.nixpkgs-fmt
-    pkgs.nodePackages.prettier
-    pkgs.rustfmt
-    pkgs.clang-tools
-  ];
-}
+  flake = import flake-compat {
+    inherit system;
+    src = ./.;
+  };
+in
+  flake.shellNix
diff --git a/treefmt.nix b/treefmt.nix
new file mode 100644
index 0000000..dadec94
--- /dev/null
+++ b/treefmt.nix
@@ -0,0 +1,76 @@
+{
+  projectRootFile = "flake.nix";
+  programs = {
+    # clang
+    clang-format.enable = true;
+
+    # deno (faster compared to prettier)
+    deno.enable = true;
+
+    # json
+    jsonfmt.enable = true;
+
+    # markdown
+    mdformat.enable = true;
+
+    # nix
+    alejandra.enable = true;
+    deadnix.enable = true;
+    statix.enable = true;
+
+    # swift
+    swift-format.enable = true;
+
+    # rust
+    rustfmt = {
+      enable = true;
+      edition = "2018";
+    };
+
+    # shell
+    shfmt.enable = true;
+
+    # yaml
+    yamlfmt.enable = true;
+  };
+  settings.formatter = {
+    # clang-format
+    clang-format = {
+      excludes = [
+        "bindings/node/binding.cc"
+        "src/parser.c"
+        "src/tree_sitter/alloc.h"
+        "src/tree_sitter/array.h"
+        "src/tree_sitter/parser.h"
+      ];
+    };
+
+    # deno
+    deno = {
+      includes = ["*.js"];
+      excludes = ["*.json"];
+    };
+
+    # jsonfmt
+    jsonfmt = {
+      excludes = ["src/**.json"];
+    };
+
+    # nix
+    statix = {
+      priority = 1;
+      excludes = ["test/*"];
+    };
+    deadnix = {
+      priority = 2;
+      excludes = ["test/*"];
+    };
+    alejandra = {
+      priority = 3;
+      excludes = ["test/*"];
+    };
+
+    # yaml
+    yamlfmt.includes = ["*.yaml" "*.yml" "*.yamlfmt"];
+  };
+}
diff --git a/treefmt.toml b/treefmt.toml
deleted file mode 100644
index e84ae51..0000000
--- a/treefmt.toml
+++ /dev/null
@@ -1,35 +0,0 @@
-[formatter.nix]
-command = "nixpkgs-fmt"
-includes = ["*.nix"]
-excludes = ["test/**.nix"]
-
-[formatter.prettier]
-command = "prettier"
-options = ["--write"]
-includes = [
-    "*.css",
-    "*.html",
-    "*.js",
-    "*.json",
-    "*.jsx",
-    "*.md",
-    "*.mdx",
-    "*.scss",
-    "*.ts",
-]
-excludes = ["src/**.json"]
-
-[formatter.rust]
-command = "rustfmt"
-options = ["--edition", "2018"]
-includes = ["*.rs"]
-
-[formatter.c]
-command = "clang-format"
-options = [ "-i" ]
-includes = [ "*.c", "*.cpp", "*.cc", "*.h", "*.hpp" ]
-excludes = [
-  "bindings/node/binding.cc",
-  "src/parser.c",
-  "src/tree_sitter/parser.h",
-]