A nice upgrade to that is to put root in a tmpfs RAM filesystem instead of ZFS:

https://elis.nu/blog/2020/05/nixos-tmpfs-as-root/

That way it doesn't even need to bother with resetting to ZFS snapshots, instead it just wipes root on shutdown and reconstructs it in RAM on reboot.

Then, optionally, with some extra work you can put /home in tmpfs too:

https://elis.nu/blog/2020/06/nixos-tmpfs-as-home/

That setup uses Home Manager, so maybe it's not for you, but worth mentioning if we're talking about making all state declarative and reproducible. You have to use the Impermanence module and set up some soft links to permanent home folders on different drive or partition. But for making all state on the system reproducible and declarative, this is the best way afaik.