Modified the code to reduce code duplication