Add missing features
Fix mistake in previous commit