diff --git a/init.el b/init.el
index d6c4457..f28b015 100644
--- a/init.el
+++ b/init.el
@@ -89,6 +89,10 @@
   :config
   (add-hook 'prog-mode-hook #'ws-butler-mode))
 
+;; Code folding
+(use-package ts-fold
+  :straight (ts-fold :type git :host github :repo "emacs-tree-sitter/ts-fold"))
+
 ;; Treesitter
 (use-package tree-sitter
   :config
@@ -96,6 +100,8 @@
     :config
     (tree-sitter-require 'rust)
     (tree-sitter-require 'java))
+  (add-hook 'tree-sitter-after-on-hook #'ts-fold-mode)
+  (add-hook 'tree-sitter-after-on-hook #'tree-sitter-hl-mode)
   (global-tree-sitter-mode +1))
 
 ;; Rust support
@@ -108,6 +114,7 @@
   (setq evil-mode-beyond-eol t)
   (setq evil-mode-fine-undo t)
   (setq evil-default-state 'emacs)
+  (setq evil-undo-system 'undo-redo)
   :config
   (evil-set-initial-state 'prog-mode 'normal)
   (evil-mode +1))
@@ -175,6 +182,7 @@
                                company-sort-prefer-same-case-prefix
                                company-sort-by-backend-importance
                                company-sort-by-occurrence))
+  (define-key company-active-map (kbd "<return>") nil)
   (define-key company-active-map (kbd "M-`") #'company-complete-selection)
   (define-key company-active-map (kbd "M-n") #'company-select-next)
   (define-key company-active-map (kbd "M-p") #'company-select-previous)