Small improvements. Actually passes a test now.