Small Changes to State logic