Refactoring: Improving your code

Modern software engineering is an iterative process. In agile methodologies, the aim is to end each sprint with a deliverable product, even if that means reducing the scope of functionality. This has the side effect that code written for the scope of this sprint, may need to be changed or replaced in a future sprint. Do this enough times, and you start to accrue what’s called “technical debt” – little hacks here and there to work around some limitation imposed by the projects architecture.

The problem with this situation is that these “little hacks” tend to accumulate, and often times they start to become not-so-little hacks. It is at this point (though realistically it should have been before this point) that the developers reach a fork in the road: Continue applying band aids with the hopes that everything will continue to be “good enough” or perform some major surgery and refactor the code base to bring the architecture back in line with scope of the project.

The decision to refactor is not one to be taken lightly. Depending on what needs to be done, a significant amount of time effort needs to be applied, and in the business world time is money. Time spent re-writing already existing code is time spent not developing new features.

In this post I will identify an area in a code base that while properly implemented for its current task is difficult to scale and proceed to refactor it into a better performing, scalable solution.

Here Be Dragons Ants…

A few posts back I covered implementing Langton’s Ants in Java. At the end of that post, I discussed adding more states for the ant to transition thru in order to observe more complex behaviors, and in here lies the rub: for every additional color state, The following changes have to be made:

  • A color representing that state needs to be added to the TileColors enumeration.
  • A new color state has to be added to the Ants move() state machine.
  • A new output color needs to be mapped to the state color and handled in the display method.

Additionally If we want to add a new type of turn-action for the ant, we need to map four new direction transitions as well. As it turns out, the turn transition table is a great place to start, as it will better exemplify what it is we need to accomplish. Right now turns are handled by using a switch statement to select the next direction based on the current direction:

private void turnRight() {
        switch(direction) {
            case NORTH:
                direction = Direction.EAST;
                break;
            case SOUTH:
                direction = Direction.WEST;
                break;
            case EAST:
                direction = Direction.SOUTH;
                break;
            case WEST:
                direction = Direction.NORTH;
                break;
        }
    }
    private void turnLeft() {
        switch(direction) {
            case NORTH:
                direction = Direction.WEST;
                break;
            case SOUTH:
                direction = Direction.EAST;
                break;
            case EAST:
                direction = Direction.NORTH;
                break;
            case WEST:
                direction = Direction.SOUTH;
                break;
        }
    }

As far as memory usage goes, this implementation is very light, unfortunately, the byte code generated for a method like this is far from sparse. Can you think of a better way to choose the next state? I’ll give you a hint: We want to replace the transisiton table shown above with a more efficient kind of table.

If you couldn’t have guessed, I’m talking about a hash table. We’re going to add two hash tables to our Ant object, leftTransition and rightTransition. These hash tables will map a direction to what the next direction should be. This will turn the two methods above into the following:

public Ant() {
rightTurns = new HashMap<>(); leftTurns = new HashMap<>();
        rightTurns.put(Direction.NORTH, Direction.EAST); leftTurns.put(Direction.NORTH, Direction.WEST);
        rightTurns.put(Direction.SOUTH, Direction.WEST); leftTurns.put(Direction.SOUTH, Direction.EAST);
        rightTurns.put(Direction.EAST, Direction.SOUTH); leftTurns.put(Direction.EAST, Direction.NORTH);
        rightTurns.put(Direction.WEST, Direction.NORTH); leftTurns.put(Direction.WEST, Direction.SOUTH);
}
private void turnRight() {
    direction = rightTurns.get(direction);
}
private void turnLeft() {
  direction = leftTurns.get(direction);
}

The boost in performance to turnLeft() and turnRight() is substantial, and now adding different kind of turns is simple as adding a new hash table with the corresponding direction transitions. Here we are going to use a different strategy. Before I had settled on using Swing for the output, I needed a way to represent tile color, so i had created the TileColor enum. This required the extra step of mapping the entry in TileColor to a coorespong Java.awt.Color. Having decided on using Swing, the TileColor enum is no longer necessary and Tile objects can now be embedded with a Color object directly, turning this:

public void paint(Graphics g) {
for (int y = 0; y < world.HEIGHT; y++) {
for (int x = 0; x < world.WIDTH; x++) {
switch (world.grid[y][x].tileColor) {
case WHITE:
g.setColor(Color.WHITE);
break;
case BLACK:
g.setColor(Color.BLACK);
break;
case GREEN:
g.setColor(Color.GREEN);
break;
case BLUE:
g.setColor(Color.BLUE);
break;
}
g.fillRect(x * ZOOM, y * ZOOM, ZOOM, ZOOM);
}
}
}

into this:

 @Override
 public void paint(Graphics g) {
      world.doStep();
      for (int y = 0; y < world.HEIGHT; y++) {
         for (int x = 0; x < world.WIDTH; x++) {
              g.setColor(world.grid[y][x].tileColor);
              g.fillRect(x * ZOOM, y * ZOOM, ZOOM, ZOOM);
         } 
    }
}

And any other place in the code that used ‘TileColor’ simply needs to be changed to ‘Color’.

Now that we have the basic strategy understood, its time to perform some deeper surgery. As implemented, the Ant.move() method is still using a clunky switch/case statement, but this case is a little trickier, because we need to also factor in method calls that are part of our state transitions. The goal is to change this:

public void move(Tile[][] grid) {
switch (grid[location.y][location.x].tileColor) {
case BLACK:
turnRight();
grid[location.y][location.x].tileColor = Color.WHITE;
break;
case WHITE:
turnLeft();
grid[location.y][location.x].tileColor = Color.GREEN;
break;
case GREEN:
turnLeft();
grid[location.y][location.x].tileColor = Color.BLUE;
break;
case BLUE:
turnLeft();
grid[location.y][location.x].tileColor = Color.BLACK;
break;
}

int nextX = (Math.abs(((location.x + compass.get(direction.ordinal()).x) + WIDTH) % WIDTH));
int nextY = (Math.abs(((location.y + compass.get(direction.ordinal()).y) + HEIGHT) % HEIGHT));

location.x = nextX;
location.y = nextY;
}

And refactor it into this:

 
public void move(Tile[][] grid) {
        changeColorAndDirectionState(grid, stateMaps.rules.get(grid[location.y][location.x].tileColor));
        takeStep(grid);
    }

The first thing i wanted to do was get rid of that switch/case statement and replace it with the HashMap technique. I did this by introducing a new object: RulePair. A rule pair is exactly what it sounds like, its an object that aggregates the two things that happen in each case statement: a left or right turn is taken, and the Tiles color is changed. So a RulePair contains a TurnDirection, and the next Tile Color:

 
public enum TurnDirection {
    LEFT,
    RIGHT;
}

public class RulePair {
    public TurnDirection turning;
    public Color nextColor;
    public RulePair(TurnDirection t, Color c) {
        turning = t;
        nextColor = c;
    }
}

Now we can create a HashMap that maps TileColor to RulePair in the same way our switch/case statement did:

     rules = new HashMap<Color, RulePair>();
rules.put(TileColor.BLACK, new RulePair(TurnDirection.RIGHT, Color.WHITE));
    rules.put(TileColor.WHITE, new RulePair(TurnDirection.LEFT, Color.GREEN));
    rules.put(TileColor.GREEN, new RulePair(TurnDirection.RIGHT, Color.BLUE));
    rules.put(TileColor.BLUE, new RulePair(TurnDirection.LEFT, Color.RED));
    rules.put(TileColor.RED, new RulePair(TurnDirection.LEFT, Color.MAGENTA));
     rules.put(TileColor.MAGENTA, new RulePair(TurnDirection.RIGHT, Color.BLACK));

And now we implement the changeColorAndDirectionState() method that consumes the HashMap entry and performs the relevant action:

    private void flipColor(Tile[][] grid, Color color) {
        grid[location.y][location.x].tileColor = color;
    }
    private void changeColorAndDirectionState(Tile[][] grid, RulePair rule) {
        if (rule.turning == TurnDirection.LEFT) {
            turnLeft();
            flipColor(grid, rule.nextColor);
        } else {
            turnRight();
            flipColor(grid, rule.nextColor);
        }
    }

And to round out the changes, the takeStep method:

 
private void takeStep(Tile[][] grid) {
        int nextX = (Math.abs(((location.x + compass.get(direction.ordinal()).x) + WIDTH) % WIDTH));
        int nextY = (Math.abs(((location.y + compass.get(direction.ordinal()).y) + HEIGHT) % HEIGHT));
        if (grid[nextY][nextX].occupied == false) {
            grid[location.y][location.x].occupied = false;
            location.x = nextX;
            location.y = nextY;
            grid[location.y][location.x].occupied = true;
        }
    }

Now that we have a working implementation, I want to go over another optimization we can refactor into the code base. Our current implementation of Langston’s ants allows us to create as many ants as the user desires. If we stay with storing the rightTurns, leftTurns, and rules HashMaps inside the Ant object, than each ant object instantiated gets their own personal memory eating copy of those three hashmaps. This is incredibly wasteful because the rules are the same for every Ant! If you’re versed in design patterns you probably already know thee answer to this conundrum: The singleton pattern. Instead of each ant getting a copy, we will create a StateMaps object to hold our three hash maps, and this object will be instantiated just once and shared amongst the ants:

 import java.util.HashMap;

public class StateMaps {
    public final HashMap<Direction, Direction> rightTurns;
    public final HashMap<Direction, Direction> leftTurns;
    public final HashMap<TileColor, RulePair> rules;
    public StateMaps() {
        rightTurns = new HashMap<>();
        leftTurns = new HashMap<>();
        rules = new HashMap<>();
        rightTurns.put(Direction.NORTH, Direction.EAST);
        rightTurns.put(Direction.SOUTH, Direction.WEST);
        rightTurns.put(Direction.EAST, Direction.SOUTH);
        rightTurns.put(Direction.WEST, Direction.NORTH);

        leftTurns.put(Direction.NORTH, Direction.WEST);
        leftTurns.put(Direction.SOUTH, Direction.EAST);
        leftTurns.put(Direction.EAST, Direction.NORTH);
        leftTurns.put(Direction.WEST, Direction.SOUTH);
       
        rules.put(TileColor.BLACK, new RulePair(TurnDirection.RIGHT, Color.WHITE));
        rules.put(TileColor.WHITE, new RulePair(TurnDirection.LEFT, Color.GREEN));
        rules.put(TileColor.GREEN, new RulePair(TurnDirection.RIGHT, Color.BLUE));
        rules.put(TileColor.BLUE, new RulePair(TurnDirection.LEFT, Color.RED));
        rules.put(TileColor.RED, new RulePair(TurnDirection.LEFT, Color.MAGENTA));
        rules.put(TileColor.MAGENTA, new RulePair(TurnDirection.RIGHT, Color.BLACK));
    }
}

We’re getting there, but anytime we want to try a new automaton with different rules, we still have to go in and make a bunch of manual changes to the rules hash table. Wouldnt it be great if we could represent rules by just passing in a text string to the StateMaps constructor? It sure would. We are going to make a trade off here, we are going to trade fine control over the order of the colors to gain the ability to easily change the rule set:

 
public StateMaps(String ruleSet) {
        rightTurns = new HashMap<>();
        leftTurns = new HashMap<>();
        rules = new HashMap<>();
        rightTurns.put(Direction.NORTH, Direction.EAST);
        rightTurns.put(Direction.SOUTH, Direction.WEST);
        rightTurns.put(Direction.EAST, Direction.SOUTH);
        rightTurns.put(Direction.WEST, Direction.NORTH);

        leftTurns.put(Direction.NORTH, Direction.WEST);
        leftTurns.put(Direction.SOUTH, Direction.EAST);
        leftTurns.put(Direction.EAST, Direction.NORTH);
        leftTurns.put(Direction.WEST, Direction.SOUTH);
   
        List<Color> colors = List.of(Color.BLACK,
                                    Color.DARK_GRAY,
                                    Color.GRAY,
                                    Color.LIGHT_GRAY,
                                    Color.WHITE,
                                    Color.CYAN,
                                    Color.BLUE,
                                    Color.GREEN,
                                    Color.MAGENTA,
                                    Color.YELLOW,
                                    Color.ORANGE,
                                    Color.PINK,
                                    Color.RED
                                    );
        int i = 0;
        int rule = 0;
        while (rule < ruleSet.length()) {
            Color curr = colors.get(i);
            i = (i+1 == colors.size()) ? 0:i+1;
            Color next = colors.get(i);
            char  direction = ruleSet.charAt(rule++);
            if (rule+1 == ruleSet.length()) {
                next = colors.get(0);
            }
            rules.put(curr, new RulePair(direction == 'L' ? TurnDirection.LEFT:TurnDirection.RIGHT, next));
        }
    }

Not only are we making our code easier to scale, it has also become much easier to read and maintain, not to mention faster and more memory efficient.

The major lesson to take away is that we shouldn’t be afraid to improve code that seemingly does not need a change a design. Could we have left this project as it was, expanding the case statements with each addition? Sure, but remember that this was a small example. On code bases that stretch to tens of thousands of lines, the impact can be significant.

Until next time friends, Happy Hacking!

Further Reading:

Refactoring: Improving the Design of Existing Code, By Martin Fowler.

The full implementation of Langton’s ants with the discussed optimizations is available at: https://github.com/maxgoren/CellularAutomata