Add the count function