Slightly improved code