Refactor state methods a bit