refactor square code
refactor square code