Add display state code