diff --git a/rust/flake.nix b/rust/flake.nix
index 34e6217..8159481 100644
--- a/rust/flake.nix
+++ b/rust/flake.nix
@@ -6,15 +6,25 @@
 
     flake-parts.url = "github:hercules-ci/flake-parts";
 
+    treefmt-nix = {
+      url = "github:numtide/treefmt-nix";
+      inputs.nixpkgs.follows = "nixpkgs";
+    };
+
     crane.url = "github:ipetkov/crane";
   };
 
   outputs = {
     flake-parts,
+    treefmt-nix,
     crane,
     ...
   } @ inputs:
     flake-parts.lib.mkFlake {inherit inputs;} {
+      imports = [
+        treefmt-nix.flakeModule
+      ];
+
       systems = ["aarch64-linux" "x86_64-linux" "aarch64-darwin" "x86_64-darwin"];
 
       perSystem = {pkgs, ...}: let
@@ -52,7 +62,13 @@
         };
 
         # Formatter for nix files, available through 'nix fmt'
-        formatter = pkgs.alejandra;
+        treefmt = {
+          programs.alejandra.enable = true;
+          programs.rustfmt = {
+            enable = true;
+            edition = (builtins.fromTOML (builtins.readFile "${src}/Cargo.toml")).package.edition;
+          };
+        };
 
         # Your custom packages
         # Accessible through 'nix build', 'nix shell', 'nix run', etc
@@ -67,10 +83,6 @@
             inherit src buildInputs nativeBuildInputs cargoArtifacts;
             cargoClippyExtraArgs = "-- --deny warnings";
           };
-
-          fmt = craneLib.cargoFmt {
-            inherit src buildInputs nativeBuildInputs;
-          };
         };
       };
     };